From f950ce1935246f1755b6de15d9270cacb24b65a4 Mon Sep 17 00:00:00 2001
From: Marvin Liu <marvin@amplitude.com>
Date: Wed, 17 Aug 2022 20:01:54 -0700
Subject: [PATCH] feat: add library context field

---
 src/main/java/com/amplitude/Amplitude.java    | 15 ++++++++
 src/main/java/com/amplitude/Event.java        |  9 +++++
 .../java/com/amplitude/AmplitudeTest.java     | 35 +++++++++++++++++++
 src/test/java/com/amplitude/EventTest.java    |  1 +
 4 files changed, 60 insertions(+)

diff --git a/src/main/java/com/amplitude/Amplitude.java b/src/main/java/com/amplitude/Amplitude.java
index 916e866..356bbfb 100644
--- a/src/main/java/com/amplitude/Amplitude.java
+++ b/src/main/java/com/amplitude/Amplitude.java
@@ -24,6 +24,7 @@ public class Amplitude {
   private Object eventQueueLock = new Object();
   private Plan plan;
   private long flushTimeout;
+  private String libraryContext;
 
   /**
    * A dictionary of key-value pairs that represent additional instructions for server save
@@ -163,6 +164,16 @@ public Amplitude setPlan(Plan plan) {
     return this;
   }
 
+  /**
+   * Set library context information.
+   *
+   * @param libraryContext String
+   */
+  public Amplitude setLibraryContext(String libraryContext) {
+    this.libraryContext = libraryContext;
+    return this;
+  }
+
   /**
    * Set custom proxy for https requests
    *
@@ -242,6 +253,10 @@ public void logEvent(Event event, AmplitudeCallbacks callbacks, MiddlewareExtra
       event.plan = this.plan;
     }
 
+    if (event.libraryContext == null) {
+      event.libraryContext = this.libraryContext;
+    }
+
     if (!middlewareRunner.run(new MiddlewarePayload(event, extra))) {
       return;
     }
diff --git a/src/main/java/com/amplitude/Event.java b/src/main/java/com/amplitude/Event.java
index 4e015f3..29f47dd 100644
--- a/src/main/java/com/amplitude/Event.java
+++ b/src/main/java/com/amplitude/Event.java
@@ -183,6 +183,11 @@ public class Event {
    */
   public Plan plan;
 
+  /**
+   * The library context information.
+   */
+  public String libraryContext;
+
   /**
    * Callback for Event
    */
@@ -279,6 +284,10 @@ public JSONObject toJsonObject() {
       if (plan != null) {
         event.put("plan", plan.toJSONObject());
       }
+
+      if (libraryContext != null) {
+        event.put("library_context", libraryContext);
+      }
     } catch (JSONException e) {
       e.printStackTrace();
     }
diff --git a/src/test/java/com/amplitude/AmplitudeTest.java b/src/test/java/com/amplitude/AmplitudeTest.java
index d6f883b..e8c2bf7 100644
--- a/src/test/java/com/amplitude/AmplitudeTest.java
+++ b/src/test/java/com/amplitude/AmplitudeTest.java
@@ -446,6 +446,41 @@ public void testSetPlan()
     assertEquals(versionId, result.getString(Constants.AMP_PLAN_VERSION_ID));
   }
 
+  @Test
+  public void testSetLibraryContext()
+          throws NoSuchFieldException, IllegalAccessException, AmplitudeInvalidAPIKeyException,
+          InterruptedException {
+    Amplitude amplitude = Amplitude.getInstance("testSetLibraryContext");
+    amplitude.init(apiKey);
+
+    String libraryContext = "test_library_context/1.x";
+
+    amplitude.setLibraryContext(libraryContext);
+
+    amplitude.useBatchMode(false);
+    amplitude.setEventUploadThreshold(1);
+    HttpCall httpCall = getMockHttpCall(amplitude, false);
+    Response response = ResponseUtil.getSuccessResponse();
+    CountDownLatch latch = new CountDownLatch(1);
+
+    AtomicReference<List<Event>> sentEvents = new AtomicReference<>();
+    when(httpCall.makeRequest(anyList()))
+            .thenAnswer(
+                    invocation -> {
+                      sentEvents.set(invocation.getArgument(0));
+                      latch.countDown();
+                      return response;
+                    });
+
+    amplitude.logEvent(new Event("test-event", "test-user"));
+
+    assertTrue(latch.await(1L, TimeUnit.SECONDS));
+    assertEquals(1, sentEvents.get().size());
+
+    Event sentEvent = sentEvents.get().get(0);
+    assertEquals(libraryContext, sentEvent.libraryContext);
+  }
+
   private HttpCall getMockHttpCall(Amplitude amplitude, boolean useBatch)
       throws NoSuchFieldException, IllegalAccessException {
     HttpCall httpCall = mock(HttpCall.class);
diff --git a/src/test/java/com/amplitude/EventTest.java b/src/test/java/com/amplitude/EventTest.java
index 94de2d5..4da8f60 100644
--- a/src/test/java/com/amplitude/EventTest.java
+++ b/src/test/java/com/amplitude/EventTest.java
@@ -52,6 +52,7 @@ public void testToJsonObject() {
     assertEquals(
         Constants.SDK_LIBRARY + "/" + Constants.SDK_VERSION, truncatedEvent.getString("library"));
     assertEquals(-1, truncatedEvent.getLong("session_id"));
+    assertFalse(truncatedEvent.has("library_context"));
   }
 
   @Test