Allow unsetting flag dimension with null

2 files changed, 11 insertions, 1 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/FetchVector.java b/flags/src/main/java/com/yahoo/vespa/flags/FetchVector.java
index 866d4782a2a..5897a5ab58b 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/FetchVector.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/FetchVector.java
@@ -84,8 +84,12 @@ public class FetchVector {
 
     public boolean isEmpty() { return map.isEmpty(); }
 
-    /** Returns a new FetchVector, identical to {@code this} except for its value in {@code dimension}. */
+    /**
+     * Returns a new FetchVector, identical to {@code this} except for its value in {@code dimension}.
+     * Dimension is removed if the value is null.
+     */
     public FetchVector with(Dimension dimension, String value) {
+        if (value == null) return makeFetchVector(merged -> merged.remove(dimension));
         return makeFetchVector(merged -> merged.put(dimension, value));
     }
 
diff --git a/flags/src/test/java/com/yahoo/vespa/flags/json/FlagDataTest.java b/flags/src/test/java/com/yahoo/vespa/flags/json/FlagDataTest.java
index 9eaa4ae4504..7f939d10bb3 100644
--- a/flags/src/test/java/com/yahoo/vespa/flags/json/FlagDataTest.java
+++ b/flags/src/test/java/com/yahoo/vespa/flags/json/FlagDataTest.java
@@ -64,6 +64,12 @@ public class FlagDataTest {
                 .with(FetchVector.Dimension.HOSTNAME, "host1")
                 .with(FetchVector.Dimension.APPLICATION_ID, "app3"));
 
+        // Verify unsetting a dimension with null works.
+        verify(Optional.of("true"), vector
+                .with(FetchVector.Dimension.HOSTNAME, "host1")
+                .with(FetchVector.Dimension.APPLICATION_ID, "app3")
+                .with(FetchVector.Dimension.APPLICATION_ID, null));
+
         // No rules apply if zone is overridden to an unknown zone
         verify(Optional.empty(), vector.with(FetchVector.Dimension.ZONE_ID, "unknown zone"));
     }