fix: add context field in the payload (#135)

* add context field in the payload * throw on context being in data
spotify · May 30, 2024 · 7f09f3e · 7f09f3e
1 parent 87407c7
commit 7f09f3e
Show file tree

Hide file tree

Showing 10 changed files with 65 additions and 16 deletions.
diff --git a/src/main/java/com/spotify/confidence/Confidence.java b/src/main/java/com/spotify/confidence/Confidence.java
@@ -101,9 +101,13 @@ public void track(String eventName) {
   }
 
   @Override
-  public void track(String eventName, ConfidenceValue.Struct message) {
+  public void track(String eventName, ConfidenceValue.Struct data) {
+    if (data.asMap().containsKey("context")) {
+      throw new Exceptions.InvalidContextInMessaageError(
+          "Field 'context' is not allowed in event's data");
+    }
     try {
-      client().emit(eventName, getContext(), Optional.of(message));
+      client().emit(eventName, getContext(), Optional.of(data));
     } catch (IllegalStateException e) {
       // swallow this exception
     }

diff --git a/src/main/java/com/spotify/confidence/EventSender.java b/src/main/java/com/spotify/confidence/EventSender.java
@@ -5,7 +5,7 @@
 
 @Beta
 public interface EventSender extends Contextual {
-  public void track(String eventName, ConfidenceValue.Struct message);
+  public void track(String eventName, ConfidenceValue.Struct data);
 
   public void track(String eventName);
 

diff --git a/src/main/java/com/spotify/confidence/EventSenderEngine.java b/src/main/java/com/spotify/confidence/EventSenderEngine.java
@@ -4,7 +4,7 @@
 import java.util.Optional;
 
 interface EventSenderEngine extends Closeable {
-  void emit(String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> message);
+  void emit(String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> data);
 
   void flush();
 }
diff --git a/src/main/java/com/spotify/confidence/EventSenderEngineImpl.java b/src/main/java/com/spotify/confidence/EventSenderEngineImpl.java
@@ -75,13 +75,13 @@ class EventSenderEngineImpl implements EventSenderEngine {
 
   @Override
   public void emit(
-      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> message) {
+      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> data) {
     if (intakeClosed) {
       log.warn("EventSenderEngine is closed, dropping event {}", name);
       return;
     }
     final Event event =
-        EventUploader.event(name, context, message).setEventTime(clock.getTimestamp()).build();
+        EventUploader.event(name, context, data).setEventTime(clock.getTimestamp()).build();
     if (estimatedMemoryConsumption.get() + event.getSerializedSize() > maxMemoryConsumption) {
       log.warn("EventSenderEngine is overloaded, dropping event {}", name);
       return;

diff --git a/src/main/java/com/spotify/confidence/EventUploader.java b/src/main/java/com/spotify/confidence/EventUploader.java
@@ -2,19 +2,23 @@
 
 import com.google.protobuf.Struct;
 import com.spotify.confidence.events.v1.Event;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
 interface EventUploader {
   static Event.Builder event(
-      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> message) {
+      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> data) {
+    final Map<String, ConfidenceValue> map = new HashMap<>(Map.of());
+    final ConfidenceValue.Struct dataStruct = data.orElse(ConfidenceValue.Struct.EMPTY);
+    map.putAll(dataStruct.asMap());
+    map.put("context", context);
+
     return Event.newBuilder()
         .setEventDefinition(EventSenderEngineImpl.EVENT_NAME_PREFIX + name)
-        .setPayload(
-            Struct.newBuilder()
-                .putAllFields(context.asProtoMap())
-                .putAllFields(message.orElse(ConfidenceValue.Struct.EMPTY).asProtoMap()));
+        .setPayload(Struct.newBuilder().putAllFields(ConfidenceValue.of(map).asProtoMap()));
   }
 
   CompletableFuture<Boolean> upload(List<Event> events);

diff --git a/src/main/java/com/spotify/confidence/Exceptions.java b/src/main/java/com/spotify/confidence/Exceptions.java
@@ -20,6 +20,12 @@ public ParseError(String message) {
     }
   }
 
+  public static class InvalidContextInMessaageError extends RuntimeException {
+    public InvalidContextInMessaageError(String message) {
+      super(message);
+    }
+  }
+
   public static class IllegalValueType extends Exception {
     public IllegalValueType(String message) {
       super(message);

diff --git a/src/test/java/com/spotify/confidence/ConfidenceContextTest.java b/src/test/java/com/spotify/confidence/ConfidenceContextTest.java
@@ -1,8 +1,10 @@
 package com.spotify.confidence;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.google.common.collect.ImmutableMap;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 
 public class ConfidenceContextTest {
@@ -11,6 +13,15 @@ public class ConfidenceContextTest {
   private final ResolverClientTestUtils.FakeFlagResolverClient fakeFlagResolverClient =
       new ResolverClientTestUtils.FakeFlagResolverClient();
 
+  @Test
+  public void testThrowInvalidContextInMessage() {
+    final Confidence root = Confidence.create(fakeEngine, fakeFlagResolverClient);
+    assertThrows(
+        Exceptions.InvalidContextInMessaageError.class,
+        () ->
+            root.track("hello", ConfidenceValue.of(Map.of("context", ConfidenceValue.NULL_VALUE))));
+  }
+
   @Test
   public void getContextContainsParentContextValues() {
     final Confidence root = Confidence.create(fakeEngine, fakeFlagResolverClient);

diff --git a/src/test/java/com/spotify/confidence/EventSenderEngineTest.java b/src/test/java/com/spotify/confidence/EventSenderEngineTest.java
@@ -110,7 +110,19 @@ public void testOverlappingKeysInPayload() throws InterruptedException {
                         "a",
                         Value.newBuilder().setNumberValue(0).build(),
                         "message",
-                        Value.newBuilder().setNumberValue(1).build()))
+                        Value.newBuilder().setNumberValue(1).build(),
+                        "context",
+                        Value.newBuilder()
+                            .setStructValue(
+                                Struct.newBuilder()
+                                    .putAllFields(
+                                        Map.of(
+                                            "a",
+                                            Value.newBuilder().setNumberValue(2).build(),
+                                            "message",
+                                            Value.newBuilder().setNumberValue(3).build()))
+                                    .build())
+                            .build()))
                 .build());
   }
 

diff --git a/src/test/java/com/spotify/confidence/FakeEventSenderEngine.java b/src/test/java/com/spotify/confidence/FakeEventSenderEngine.java
@@ -25,8 +25,8 @@ public void close() throws IOException {
 
   @Override
   public void emit(
-      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> message) {
-    events.add(event(name, context, message).setEventTime(clock.getTimestamp()).build());
+      String name, ConfidenceValue.Struct context, Optional<ConfidenceValue.Struct> data) {
+    events.add(event(name, context, data).setEventTime(clock.getTimestamp()).build());
   }
 
   @Override

diff --git a/src/test/java/com/spotify/confidence/GrpcEventUploaderTest.java b/src/test/java/com/spotify/confidence/GrpcEventUploaderTest.java
@@ -98,7 +98,13 @@ public void testMapsSingleEventBatchToProtobuf() throws ExecutionException, Inte
 
     final Map<String, com.google.protobuf.Value> fieldsMap = protoEvent.getPayload().getFieldsMap();
     assertThat(fieldsMap.get("messageKey").getStringValue()).isEqualTo("value_1");
-    assertThat(fieldsMap.get("contextKey").getStringValue()).isEqualTo("value_1");
+    assertThat(
+            fieldsMap
+                .get("context")
+                .getStructValue()
+                .getFieldsOrThrow("contextKey")
+                .getStringValue())
+        .isEqualTo("value_1");
   }
 
   @Test
@@ -130,7 +136,13 @@ public void testMapsMultiEventBatchToProtobuf() {
       final Map<String, com.google.protobuf.Value> fieldsMap =
           protoEvent.getPayload().getFieldsMap();
       assertThat(fieldsMap.get("messageKey").getStringValue()).isEqualTo("value_m" + (i + 1));
-      assertThat(fieldsMap.get("contextKey").getStringValue()).isEqualTo("value_c" + (i + 1));
+      assertThat(
+              fieldsMap
+                  .get("context")
+                  .getStructValue()
+                  .getFieldsOrThrow("contextKey")
+                  .getStringValue())
+          .isEqualTo("value_c" + (i + 1));
     }
   }