feat: Record exposures under local evaluation (#251)

Part of Flagsmith/flagsmith-private#325. Java counterpart of
Flagsmith/flagsmith-python-client#276, Flagsmith/flagsmith-go-client#246
and Flagsmith/flagsmith-nodejs-client#318.

`getExperimentFlag` only recorded exposures under remote evaluation. The
environment document carries `metadata.experiment` on feature states,
but local evaluation dropped it, along with the variant.

## Changes

- Map the multivariate option `key` into the evaluation context, and
carry `variant` from the evaluation result onto `Flag`. Locally
evaluated identity flags now report their variant.
- Map the feature state's `metadata.experiment` into the context's
feature metadata, and carry it onto identity flags.
- Environment flags carry no experiment: it is removed once, when the
environment evaluation context is built, so `getEnvironmentFlags` does
no extra work per call.
- `getExperimentFlag` records no exposure when the flag has no variant,
as the other SDKs do.

## Tests

- Local evaluation: an enrolled identity records one exposure with the
variant and `experiment_id`; outside the rollout, an identity override,
and a feature without an experiment record none.
- Environment flags carry no experiment under local evaluation.
- Remote: an enrolled flag without a variant records no exposure.

Author: gagantrivedi

src/main/java/com/flagsmith/FlagsmithClient.java

@@ -9,6 +9,9 @@
 import com.flagsmith.flagengine.Engine;
 import com.flagsmith.flagengine.EvaluationContext;
 import com.flagsmith.flagengine.EvaluationResult;
+import com.flagsmith.flagengine.FeatureContext;
+import com.flagsmith.flagengine.Features;
+import com.flagsmith.flagengine.SegmentContext;
 import com.flagsmith.flagengine.SegmentRule;
 import com.flagsmith.flagengine.Segments;
 import com.flagsmith.interfaces.FlagsmithCache;
@@ -79,6 +82,7 @@ public void setEvaluationContext(EvaluationContext evaluationContext) {
     this.environmentEvaluationContext = evaluationContext == null
         ? null
         : new EvaluationContext(evaluationContext)
+            .withFeatures(withoutExperiments(evaluationContext.getFeatures()))
             .withSegments(getIdentityFreeSegments(evaluationContext.getSegments()));
     this.evaluationContext = evaluationContext;
   }
@@ -250,8 +254,7 @@ public BaseFlag getExperimentFlag(String featureName, String identifier)
 
   /**
    * Get an identity's flag, recording one {@code $flag_exposure} event if the identity is enrolled
-   * in a running experiment on it. Only remote evaluation carries experiment metadata, so local
-   * evaluation and offline mode record no exposure.
+   * in a running experiment on it.
    *
    * @param featureName feature name
    * @param identifier  identifier string
@@ -298,6 +301,12 @@ public BaseFlag getExperimentFlag(
       return flag;
     }
 
+    if (((Flag) flag).getVariant() == null) {
+      logger.info("Not recording an exposure for feature {}: the flag has no variant.",
+          featureName);
+      return flag;
+    }
+
     Map<String, Object> metadata = new HashMap<>();
     metadata.put("experiment_id", experiment.getId());
     trackExposureEvent(featureName, identifier, ((Flag) flag).getVariant(), traits, metadata);
@@ -463,12 +472,40 @@ private static Segments getIdentityFreeSegments(Segments segments) {
     Segments identityFreeSegments = new Segments();
     segments.getAdditionalProperties().forEach((segmentKey, segmentContext) -> {
       if (isIdentityFree(segmentContext.getRules())) {
-        identityFreeSegments.setAdditionalProperty(segmentKey, segmentContext);
+        identityFreeSegments.setAdditionalProperty(segmentKey, segmentContext.getOverrides() == null
+            ? segmentContext
+            : new SegmentContext(segmentContext).withOverrides(segmentContext.getOverrides()
+                .stream()
+                .map(FlagsmithClient::withoutExperiment)
+                .collect(Collectors.toList())));
       }
     });
     return identityFreeSegments;
   }
 
+  /*
+   * Only an identity is enrolled in an experiment, so environment flags carry none.
+   */
+  private static Features withoutExperiments(Features features) {
+    if (features == null) {
+      return null;
+    }
+    Features environmentFeatures = new Features();
+    features.getAdditionalProperties().forEach((featureName, featureContext) ->
+        environmentFeatures.setAdditionalProperty(featureName, withoutExperiment(featureContext)));
+    return environmentFeatures;
+  }
+
+  private static FeatureContext withoutExperiment(FeatureContext featureContext) {
+    Map<String, Object> metadata = featureContext.getMetadata();
+    if (metadata == null || !metadata.containsKey("experiment")) {
+      return featureContext;
+    }
+    Map<String, Object> environmentMetadata = new HashMap<>(metadata);
+    environmentMetadata.remove("experiment");
+    return new FeatureContext(featureContext).withMetadata(environmentMetadata);
+  }
+
   private static boolean isIdentityFree(List<SegmentRule> rules) {
     return rules == null || rules.stream().allMatch(rule ->
         (rule.getConditions() == null || rule.getConditions().stream().allMatch(condition ->

src/main/java/com/flagsmith/mappers/EngineMappers.java

@@ -63,6 +63,8 @@ public static Flag mapFlagResultToFlag(
     flag.setValue(flagResult.getValue());
     flag.setEnabled(flagResult.getEnabled());
     flag.setReason(flagResult.getReason());
+    flag.setVariant(flagResult.getVariant());
+    flag.setExperiment(metadata.getExperiment());
     return flag;
   }
 
@@ -358,6 +360,7 @@ private static FeatureContext mapFeatureStateToFeatureContext(FeatureStateModel
     for (MultivariateFeatureStateValueModel mvValue :
         featureState.getMultivariateFeatureStateValues()) {
       FeatureValue variant = new FeatureValue()
+          .withKey(mvValue.getMultivariateFeatureOption().getKey())
           .withValue(mvValue.getMultivariateFeatureOption().getValue())
           .withWeight(mvValue.getPercentageAllocation().doubleValue())
           .withPriority(getMultivariateFeatureValuePriority(mvValue));
@@ -381,6 +384,9 @@ private static Map<String, Object> mapFeatureStateToFeatureMetadata(
       FeatureStateModel featureState) {
     FeatureMetadata metadata = new FeatureMetadata();
     metadata.setFlagsmithId(featureState.getFeature().getId());
+    if (featureState.getMetadata() != null) {
+      metadata.setExperiment(featureState.getMetadata().getExperiment());
+    }
     return MapperFactory.getMapper().convertValue(metadata,
         new com.fasterxml.jackson.core.type.TypeReference<Map<String, Object>>() {
         });

src/main/java/com/flagsmith/models/ExperimentMetadata.java

@@ -4,7 +4,7 @@
 import lombok.Data;
 
 /**
- * Details of the experiment running on a feature, as returned by remote evaluation.
+ * Details of the experiment running on a feature. Only identity evaluation carries it.
  */
 @Data
 public class ExperimentMetadata {

src/main/java/com/flagsmith/models/FeatureMetadata.java

@@ -1,5 +1,7 @@
 package com.flagsmith.models;
 
+import com.fasterxml.jackson.annotation.JsonInclude;
+
 /**
  * FeatureMetadata
  * 
@@ -8,6 +10,8 @@
  */
 public class FeatureMetadata {
   private Integer flagsmithId;
+  @JsonInclude(JsonInclude.Include.NON_NULL)
+  private ExperimentMetadata experiment;
 
   /*
    * FlagsmithId
@@ -28,4 +32,18 @@ public Integer getFlagsmithId() {
   public void setFlagsmithId(Integer flagsmithId) {
     this.flagsmithId = flagsmithId;
   }
+
+  /*
+   * Experiment
+   * <p>The running experiment the feature state belongs to, if any.
+   *
+   * @return The experiment
+   */
+  public ExperimentMetadata getExperiment() {
+    return experiment;
+  }
+
+  public void setExperiment(ExperimentMetadata experiment) {
+    this.experiment = experiment;
+  }
 }

src/main/java/com/flagsmith/models/features/MultivariateFeatureOptionModel.java

@@ -6,4 +6,5 @@
 @Data
 public class MultivariateFeatureOptionModel extends BaseModel {
   private Object value;
+  private String key;
 }

src/test/java/com/flagsmith/FlagsmithClientTest.java

@@ -57,6 +57,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
@@ -1148,6 +1149,10 @@ public void testGetExperimentFlagRecordsNoExposureWhenNotEnrolledOrDisabled()
         // enrolled, but the flag is off
         assertEquals(Boolean.FALSE,
                 client.getExperimentFlag("disabled_feature", "user-1").getEnabled());
+        // enrolled, but the flag has no variant
+        Flag noVariant = (Flag) client.getExperimentFlag("no_variant", "user-1");
+        assertEquals(Boolean.TRUE, noVariant.getExperiment().getInExperiment());
+        assertNull(noVariant.getVariant());
 
         verify(processor, never()).trackExposureEvent(any(), any(), any(), any(), any());
     }
@@ -1175,39 +1180,90 @@ public void testGetExperimentFlagFallsBackToTheDefaultHandlerWithoutAnExposure(
         verify(processor, never()).trackExposureEvent(any(), any(), any(), any(), any());
     }
 
-    @Test
-    public void testGetExperimentFlagRecordsNoExposureWithLocalEvaluation()
-            throws JsonProcessingException, FlagsmithClientError {
-        String baseUrl = "http://bad-url";
+    /**
+     * A locally evaluating client serving
+     * {@link FlagsmithTestHelper#getEnvironmentDocumentWithExperiment()}.
+     */
+    private static FlagsmithClient localExperimentClient(EventProcessor processor) {
         MockInterceptor interceptor = new MockInterceptor();
-        EventProcessor processor = mock(EventProcessor.class);
         interceptor.addRule()
-                .get(baseUrl + "/environment-document/")
+                .get("http://bad-url/environment-document/")
                 .anyTimes()
-                .respond(
-                        MapperFactory.getMapper()
-                                .writeValueAsString(FlagsmithTestHelper.environmentModel()),
-                        MEDIATYPE_JSON);
-
+                .respond(FlagsmithTestHelper.getEnvironmentDocumentWithExperiment(), MEDIATYPE_JSON);
         FlagsmithClient client = FlagsmithClient.newBuilder()
                 .withConfiguration(FlagsmithConfig.newBuilder()
-                        .baseUri(baseUrl)
+                        .baseUri("http://bad-url")
                         .addHttpInterceptor(interceptor)
                         .withEventProcessor(processor)
                         .withLocalEvaluation(true)
                         .build())
                 .setApiKey("ser.abcdefg")
                 .build();
         client.updateEnvironment();
+        return client;
+    }
+
+    @ParameterizedTest
+    @CsvSource({
+        "user-1, new checkout, treatment",
+        "user-3, control, control"
+    })
+    public void testGetExperimentFlagRecordsAnExposureWithLocalEvaluationWhenEnrolled(
+            String identifier, String expectedValue, String expectedVariant)
+            throws FlagsmithClientError {
+        EventProcessor processor = mock(EventProcessor.class);
+        FlagsmithClient client = localExperimentClient(processor);
+        Map<String, Object> traits = Collections.singletonMap("plan", "premium");
 
-        BaseFlag flag = client.getExperimentFlag("some_feature", "user-1");
+        Flag flag = (Flag) client.getExperimentFlag("checkout", identifier, traits);
 
-        assertTrue(flag instanceof Flag);
-        assertNull(((Flag) flag).getVariant());
-        assertNull(((Flag) flag).getExperiment());
+        assertEquals(expectedValue, flag.getValue());
+        assertEquals(expectedVariant, flag.getVariant());
+        assertEquals(42, flag.getExperiment().getId());
+        assertEquals("Checkout experiment", flag.getExperiment().getName());
+        assertEquals(Boolean.TRUE, flag.getExperiment().getInExperiment());
+        verify(processor).trackExposureEvent("checkout", identifier, expectedVariant, traits,
+                Collections.singletonMap("experiment_id", 42));
+    }
+
+    private static Stream<Arguments> localExperimentFlagsWithoutAnExposure() {
+        return Stream.of(
+                // outside the rollout
+                Arguments.of("checkout", "user-1", "free", "control", Boolean.FALSE),
+                // an identity override carries no experiment
+                Arguments.of("checkout", "user_overridden", "premium", "overridden checkout", null),
+                // no experiment on the feature
+                Arguments.of("plain", "user-1", "premium", "hello", null));
+    }
+
+    @ParameterizedTest
+    @MethodSource("localExperimentFlagsWithoutAnExposure")
+    public void testGetExperimentFlagRecordsNoExposureWithLocalEvaluationWhenNotEnrolled(
+            String featureName, String identifier, String plan, String expectedValue,
+            Boolean expectedInExperiment) throws FlagsmithClientError {
+        EventProcessor processor = mock(EventProcessor.class);
+        FlagsmithClient client = localExperimentClient(processor);
+
+        Flag flag = (Flag) client.getExperimentFlag(
+                featureName, identifier, Collections.singletonMap("plan", plan));
+
+        assertEquals(expectedValue, flag.getValue());
+        assertEquals(expectedInExperiment,
+                flag.getExperiment() == null ? null : flag.getExperiment().getInExperiment());
         verify(processor, never()).trackExposureEvent(any(), any(), any(), any(), any());
     }
 
+    @Test
+    public void testEnvironmentFlagsCarryNoExperimentWithLocalEvaluation()
+            throws FlagsmithClientError {
+        FlagsmithClient client = localExperimentClient(mock(EventProcessor.class));
+
+        Flag flag = (Flag) client.getEnvironmentFlags().getFlag("checkout");
+
+        assertEquals("control", flag.getValue());
+        assertNull(flag.getExperiment());
+    }
+
     /**
      * An events-enabled config serving the experiment identity flags, which records each events
      * batch as "environment key|body".

src/test/java/com/flagsmith/FlagsmithTestHelper.java

@@ -447,10 +447,50 @@ public static String getIdentitiesFlagsWithExperiment() {
         + "  {\"id\": 4, \"feature\": {\"id\": 4, \"name\": \"disabled_feature\", \"type\": \"MULTIVARIATE\"},\n"
         + "   \"feature_state_value\": \"off\", \"enabled\": false, \"variant\": \"treatment\",\n"
         + "   \"reason\": \"SPLIT; weight=50.0\", \"metadata\": {\"experiment\":\n"
-        + "   {\"id\": 44, \"name\": \"disabled_experiment\", \"in_experiment\": true}}}\n"
+        + "   {\"id\": 44, \"name\": \"disabled_experiment\", \"in_experiment\": true}}},\n"
+        + "  {\"id\": 5, \"feature\": {\"id\": 5, \"name\": \"no_variant\", \"type\": \"STANDARD\"},\n"
+        + "   \"feature_state_value\": \"on\", \"enabled\": true, \"metadata\": {\"experiment\":\n"
+        + "   {\"id\": 45, \"name\": \"no_variant_experiment\", \"in_experiment\": true}}}\n"
         + "]}";
   }
 
+  /**
+   * An environment document running an experiment on {@code checkout}. Its rollout segment is
+   * identities on the premium plan, split evenly between {@code control} and {@code treatment}.
+   * {@code user_overridden} has an identity override for {@code checkout}, and {@code plain} has no
+   * experiment.
+   */
+  public static String getEnvironmentDocumentWithExperiment() {
+    String experiment = "{\"id\": 42, \"name\": \"Checkout experiment\", \"in_experiment\": %s}";
+    String checkout = "{\"name\": \"checkout\", \"type\": \"MULTIVARIATE\", \"id\": 10}";
+    return "{\"api_key\": \"B62qaMZNwfiqT76p38ggrQ\", \"name\": \"Test Environment\", \"id\": 1,\n"
+        + " \"project\": {\"name\": \"Test project\", \"id\": 1, \"hide_disabled_flags\": false,\n"
+        + "  \"organisation\": {\"id\": 1, \"name\": \"Test Org\", \"feature_analytics\": false,\n"
+        + "   \"persist_trait_data\": true, \"stop_serving_flags\": false},\n"
+        + "  \"segments\": [{\"id\": 7, \"name\": \"Checkout experiment rollout\",\n"
+        + "   \"rules\": [{\"type\": \"ALL\", \"rules\": [], \"conditions\": [\n"
+        + "    {\"operator\": \"EQUAL\", \"property_\": \"plan\", \"value\": \"premium\"}]}],\n"
+        + "   \"feature_states\": [{\"id\": 2, \"featurestate_uuid\": \"c2c2c2c2-c2c2-c2c2-c2c2-c2c2c2c2c2c2\",\n"
+        + "    \"feature\": " + checkout + ", \"feature_state_value\": \"control\", \"enabled\": true,\n"
+        + "    \"feature_segment\": {\"priority\": 0},\n"
+        + "    \"multivariate_feature_state_values\": [{\"id\": 3, \"percentage_allocation\": 50,\n"
+        + "     \"mv_fs_value_uuid\": \"3e3e3e3e-3e3e-3e3e-3e3e-3e3e3e3e3e3e\",\n"
+        + "     \"multivariate_feature_option\": {\"id\": 3, \"value\": \"new checkout\", \"key\": \"treatment\"}}],\n"
+        + "    \"metadata\": {\"experiment\": " + String.format(experiment, "true") + "}}]}]},\n"
+        + " \"feature_states\": [\n"
+        + "  {\"id\": 1, \"featurestate_uuid\": \"c1c1c1c1-c1c1-c1c1-c1c1-c1c1c1c1c1c1\",\n"
+        + "   \"feature\": " + checkout + ", \"feature_state_value\": \"control\", \"enabled\": true,\n"
+        + "   \"multivariate_feature_state_values\": [],\n"
+        + "   \"metadata\": {\"experiment\": " + String.format(experiment, "false") + "}},\n"
+        + "  {\"id\": 3, \"featurestate_uuid\": \"c3c3c3c3-c3c3-c3c3-c3c3-c3c3c3c3c3c3\",\n"
+        + "   \"feature\": {\"name\": \"plain\", \"type\": \"STANDARD\", \"id\": 11},\n"
+        + "   \"feature_state_value\": \"hello\", \"enabled\": true, \"multivariate_feature_state_values\": []}],\n"
+        + " \"identity_overrides\": [{\"identifier\": \"user_overridden\", \"identity_features\": [\n"
+        + "  {\"featurestate_uuid\": \"c4c4c4c4-c4c4-c4c4-c4c4-c4c4c4c4c4c4\", \"feature\": " + checkout + ",\n"
+        + "   \"feature_state_value\": \"overridden checkout\", \"enabled\": true,\n"
+        + "   \"multivariate_feature_state_values\": []}]}]}";
+  }
+
   public static <T> Future<T> futurableReturn(T response) {
     CompletableFuture<T> promise = new CompletableFuture<>();
     promise.complete(response);