Updated TestFeatureStore for multivariate flags

src/main/java/com/launchdarkly/client/TestFeatureStore.java

@@ -20,20 +20,38 @@ public class TestFeatureStore extends InMemoryFeatureStore {
 
   private AtomicInteger version = new AtomicInteger(0);
 
+  /**
+   * Sets the value of a boolean feature flag for all users.
+   *
+   * @param key the key of the feature flag
+   * @param value the new value of the feature flag
+   */
+  public void setBooleanValue(String key, Boolean value) {
+    FeatureFlag newFeature = new FeatureFlagBuilder(key)
+            .on(false)
+            .offVariation(value ? 0 : 1)
+            .variations(TRUE_FALSE_VARIATIONS)
+            .version(version.incrementAndGet())
+            .build();
+    upsert(key, newFeature);
+  }
+
   /**
    * Turns a feature, identified by key, to evaluate to true for every user. If the feature rules already exist in the store then it will override it to be true for every {@link LDUser}.
    * If the feature rule is not currently in the store, it will create one that is true for every {@link LDUser}.
    *
    * @param key the key of the feature flag to evaluate to true.
    */
   public void setFeatureTrue(String key) {
-    FeatureFlag newFeature = new FeatureFlagBuilder(key)
-        .on(false)
-        .offVariation(0)
-        .variations(TRUE_FALSE_VARIATIONS)
-        .version(version.incrementAndGet())
-        .build();
-    upsert(key, newFeature);
+    setBooleanValue(key, true);
+  }
+
+  /**
+   * @deprecated use {@link #setFeatureTrue(String key)}
+   */
+  @Deprecated
+  public void turnFeatureOn(String key) {
+    setFeatureTrue(key);
   }
 
   /**
@@ -43,12 +61,56 @@ public void setFeatureTrue(String key) {
    * @param key the key of the feature flag to evaluate to false.
    */
   public void setFeatureFalse(String key) {
+    setBooleanValue(key, false);
+  }
+
+  /**
+   * @deprecated use {@link #setFeatureFalse(String key)}
+   */
+  @Deprecated
+  public void turnFeatureOff(String key) {
+    setFeatureFalse(key);
+  }
+
+  /**
+   * Sets the value of an integer multivariate feature flag, for all users.
+   * @param key the key of the flag
+   * @param value the new value of the flag
+     */
+  public void setIntegerValue(String key, Integer value) {
+    setJsonValue(key, new JsonPrimitive(value));
+  }
+
+  /**
+   * Sets the value of a double multivariate feature flag, for all users.
+   * @param key the key of the flag
+   * @param value the new value of the flag
+     */
+  public void setDoubleValue(String key, Double value) {
+    setJsonValue(key, new JsonPrimitive(value));
+  }
+
+  /**
+   * Sets the value of a string multivariate feature flag, for all users.
+   * @param key the key of the flag
+   * @param value the new value of the flag
+     */
+  public void setStringValue(String key, String value) {
+    setJsonValue(key, new JsonPrimitive(value));
+  }
+
+  /**
+   * Sets the value of a JsonElement multivariate feature flag, for all users.
+   * @param key the key of the flag
+   * @param value the new value of the flag
+     */
+  public void setJsonValue(String key, JsonElement value) {
     FeatureFlag newFeature = new FeatureFlagBuilder(key)
-        .on(false)
-        .offVariation(1)
-        .variations(TRUE_FALSE_VARIATIONS)
-        .version(version.incrementAndGet())
-        .build();
+            .on(false)
+            .offVariation(0)
+            .variations(Arrays.asList(value))
+            .version(version.incrementAndGet())
+            .build();
     upsert(key, newFeature);
   }
 }

src/test/java/com/launchdarkly/client/LDClientTest.java

@@ -1,5 +1,7 @@
 package com.launchdarkly.client;
 
+import com.google.gson.JsonArray;
+import com.google.gson.JsonPrimitive;
 import org.easymock.EasyMockSupport;
 import org.junit.Before;
 import org.junit.Test;
@@ -117,6 +119,142 @@ public void testTestFeatureStoreFlagTrueThenFalse() throws IOException, Interrup
     verifyAll();
   }
 
+  @Test
+  public void testTestFeatureStoreIntegerVariation() throws Exception {
+    TestFeatureStore testFeatureStore = new TestFeatureStore();
+    LDConfig config = new LDConfig.Builder()
+            .startWaitMillis(10L)
+            .stream(false)
+            .featureStore(testFeatureStore)
+            .build();
+
+    expect(initFuture.get(10L, TimeUnit.MILLISECONDS)).andReturn(new Object());
+    expect(pollingProcessor.start()).andReturn(initFuture);
+    expect(pollingProcessor.initialized()).andReturn(true).times(2);
+    expect(eventProcessor.sendEvent(anyObject(Event.class))).andReturn(true).times(2);
+    replayAll();
+
+    client = createMockClient(config);
+
+    testFeatureStore.setIntegerValue("key", 1);
+    assertEquals(new Integer(1), client.intVariation("key", new LDUser("user"), 0));
+    testFeatureStore.setIntegerValue("key", 42);
+    assertEquals(new Integer(42), client.intVariation("key", new LDUser("user"), 1));
+    verifyAll();
+  }
+
+  @Test
+  public void testTestFeatureStoreDoubleVariation() throws Exception {
+    TestFeatureStore testFeatureStore = new TestFeatureStore();
+    LDConfig config = new LDConfig.Builder()
+            .startWaitMillis(10L)
+            .stream(false)
+            .featureStore(testFeatureStore)
+            .build();
+
+    expect(initFuture.get(10L, TimeUnit.MILLISECONDS)).andReturn(new Object());
+    expect(pollingProcessor.start()).andReturn(initFuture);
+    expect(pollingProcessor.initialized()).andReturn(true).times(2);
+    expect(eventProcessor.sendEvent(anyObject(Event.class))).andReturn(true).times(2);
+    replayAll();
+
+    client = createMockClient(config);
+
+    testFeatureStore.setDoubleValue("key", 1d);
+    assertEquals(new Double(1), client.doubleVariation("key", new LDUser("user"), 0d));
+    testFeatureStore.setDoubleValue("key", 42d);
+    assertEquals(new Double(42), client.doubleVariation("key", new LDUser("user"), 1d));
+    verifyAll();
+  }
+
+  @Test
+  public void testTestFeatureStoreStringVariation() throws Exception {
+    TestFeatureStore testFeatureStore = new TestFeatureStore();
+    LDConfig config = new LDConfig.Builder()
+            .startWaitMillis(10L)
+            .stream(false)
+            .featureStore(testFeatureStore)
+            .build();
+
+    expect(initFuture.get(10L, TimeUnit.MILLISECONDS)).andReturn(new Object());
+    expect(pollingProcessor.start()).andReturn(initFuture);
+    expect(pollingProcessor.initialized()).andReturn(true).times(2);
+    expect(eventProcessor.sendEvent(anyObject(Event.class))).andReturn(true).times(2);
+    replayAll();
+
+    client = createMockClient(config);
+
+    testFeatureStore.setStringValue("key", "apples");
+    assertEquals("apples", client.stringVariation("key", new LDUser("user"), "oranges"));
+    testFeatureStore.setStringValue("key", "bananas");
+    assertEquals("bananas", client.stringVariation("key", new LDUser("user"), "apples"));
+    verifyAll();
+  }
+
+  @Test
+  public void testTestFeatureStoreJsonVariationPrimitive() throws Exception {
+    TestFeatureStore testFeatureStore = new TestFeatureStore();
+    LDConfig config = new LDConfig.Builder()
+            .startWaitMillis(10L)
+            .stream(false)
+            .featureStore(testFeatureStore)
+            .build();
+
+    expect(initFuture.get(10L, TimeUnit.MILLISECONDS)).andReturn(new Object());
+    expect(pollingProcessor.start()).andReturn(initFuture);
+    expect(pollingProcessor.initialized()).andReturn(true).times(4);
+    expect(eventProcessor.sendEvent(anyObject(Event.class))).andReturn(true).times(4);
+    replayAll();
+
+    client = createMockClient(config);
+
+    // Character
+    testFeatureStore.setJsonValue("key", new JsonPrimitive('a'));
+    assertEquals(new JsonPrimitive('a'), client.jsonVariation("key", new LDUser("user"), new JsonPrimitive('b')));
+    testFeatureStore.setJsonValue("key", new JsonPrimitive('b'));
+    assertEquals(new JsonPrimitive('b'), client.jsonVariation("key", new LDUser("user"), new JsonPrimitive('z')));
+
+    // Long
+    testFeatureStore.setJsonValue("key", new JsonPrimitive(1L));
+    assertEquals(new JsonPrimitive(1l), client.jsonVariation("key", new LDUser("user"), new JsonPrimitive(0L)));
+    testFeatureStore.setJsonValue("key", new JsonPrimitive(42L));
+    assertEquals(new JsonPrimitive(42L), client.jsonVariation("key", new LDUser("user"), new JsonPrimitive(0L)));
+    verifyAll();
+  }
+
+  @Test
+  public void testTestFeatureStoreJsonVariationArray() throws Exception {
+    TestFeatureStore testFeatureStore = new TestFeatureStore();
+    LDConfig config = new LDConfig.Builder()
+            .startWaitMillis(10L)
+            .stream(false)
+            .featureStore(testFeatureStore)
+            .build();
+
+    expect(initFuture.get(10L, TimeUnit.MILLISECONDS)).andReturn(new Object());
+    expect(pollingProcessor.start()).andReturn(initFuture);
+    expect(pollingProcessor.initialized()).andReturn(true).times(2);
+    expect(eventProcessor.sendEvent(anyObject(Event.class))).andReturn(true).times(2);
+    replayAll();
+
+    client = createMockClient(config);
+
+    // JsonArray
+    JsonArray array = new JsonArray();
+    array.add("red");
+    array.add("blue");
+    array.add("green");
+    testFeatureStore.setJsonValue("key", array);
+    assertEquals(array, client.jsonVariation("key", new LDUser("user"), new JsonArray()));
+
+    JsonArray array2 = new JsonArray();
+    array2.addAll(array);
+    array2.add("yellow");
+    testFeatureStore.setJsonValue("key", array2);
+    assertEquals(array2, client.jsonVariation("key", new LDUser("user"), new JsonArray()));
+    verifyAll();
+  }
+
   @Test
   public void testUseLdd() throws IOException {
     LDConfig config = new LDConfig.Builder()