core: add remove method for SessionContext (Netflix#1087)

argha-c · Mar 18, 2022 · 2e05ce2 · 2e05ce2
1 parent bc68a1a
commit 2e05ce2
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/zuul-core/src/main/java/com/netflix/zuul/context/SessionContext.java b/zuul-core/src/main/java/com/netflix/zuul/context/SessionContext.java
@@ -196,6 +196,23 @@ public <T> boolean remove(Key<T> key, T value) {
         return res;
     }
 
+    /**
+     * {@inheritDoc}
+     *
+     * <p>This method exists for static analysis.
+     */
+    @Override
+    public Object remove(Object key) {
+        return super.remove(key);
+    }
+
+    public <T> T remove(Key<T> key) {
+        Objects.requireNonNull(key, "key");
+        @SuppressWarnings("unchecked") // sorry
+        T res = ((Map<Key<T>, T>) (Map) typedMap).remove(key);
+        return res;
+    }
+
     public Set<Key<?>> keys() {
         return Collections.unmodifiableSet(new HashSet<>(typedMap.keySet()));
     }
@@ -206,6 +223,7 @@ public Set<Key<?>> keys() {
     @Override
     public SessionContext clone()
     {
+        // TODO(carl-mastrangelo): copy over the type safe keys
         return (SessionContext) super.clone();
     }
 

diff --git a/zuul-core/src/test/java/com/netflix/zuul/context/SessionContextTest.java b/zuul-core/src/test/java/com/netflix/zuul/context/SessionContextTest.java
@@ -94,4 +94,17 @@ public void getOrDefault_failsOnNullDefault() {
 
         assertThrows(NullPointerException.class, () -> context.getOrDefault(key, null));
     }
+
+    @Test
+    public void remove() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+        context.put(key, "bar");
+
+        Truth.assertThat(context.get(key)).isEqualTo("bar");
+
+        String val = context.remove(key);
+        Truth.assertThat(context.get(key)).isNull();
+        Truth.assertThat(val).isEqualTo("bar");
+    }
 }