Code review fixes

2 files changed, 11 insertions, 3 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
index f2d96de79dc..cc4d55c0c94 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
@@ -8,6 +8,7 @@ import java.util.List;
 import java.util.Optional;
 import java.util.TreeMap;
 
+import static com.yahoo.vespa.flags.FetchVector.Dimension.APPLICATION_ID;
 import static com.yahoo.vespa.flags.FetchVector.Dimension.HOSTNAME;
 
 /**
@@ -41,7 +42,7 @@ public class Flags {
             "use-config-server-cache", true,
             "Whether config server will use cache to answer config requests.",
             "Takes effect immediately when changed.",
-            HOSTNAME, FetchVector.Dimension.APPLICATION_ID);
+            HOSTNAME, APPLICATION_ID);
 
     public static final UnboundBooleanFlag CONFIG_SERVER_BOOTSTRAP_IN_SEPARATE_THREAD = defineFeatureFlag(
             "config-server-bootstrap-in-separate-thread", true,
diff --git a/flags/src/main/java/com/yahoo/vespa/flags/InMemoryFlagSource.java b/flags/src/main/java/com/yahoo/vespa/flags/InMemoryFlagSource.java
index 25b871f8cd1..5e2afe39dd8 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/InMemoryFlagSource.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/InMemoryFlagSource.java
@@ -1,6 +1,8 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.flags;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.DoubleNode;
 import com.yahoo.vespa.flags.json.FlagData;
 import com.yahoo.vespa.flags.json.Rule;
 
@@ -21,8 +23,13 @@ public class InMemoryFlagSource implements FlagSource {
         return this;
     }
 
-    public InMemoryFlagSource withFlag(FlagId flagId, FetchVector defaultFetchVector, Rule... rules) {
-        flagDataById.put(flagId, new FlagData(flagId, defaultFetchVector, rules));
+    public InMemoryFlagSource withDoubleFlag(FlagId flagId, double value) {
+        return withRawFlag(flagId, new DoubleNode(value));
+    }
+
+    private InMemoryFlagSource withRawFlag(FlagId flagId, JsonNode rawFlag) {
+        Rule rule = new Rule(Optional.of(JsonNodeRawFlag.fromJsonNode(rawFlag)));
+        flagDataById.put(flagId, new FlagData(flagId, new FetchVector(), rule));
         return this;
     }