From 027099bd808109c5619fb15e9a5cfe2608f3dff2 Mon Sep 17 00:00:00 2001
From: Connor Worley <cworley@netflix.com>
Date: Sun, 3 Nov 2024 19:34:56 -0800
Subject: [PATCH] Add Key<T> specialization to SessionContext.containsKey

---
 .../netflix/zuul/context/SessionContext.java  | 17 ++++++++++++++++
 .../zuul/context/SessionContextTest.java      | 20 +++++++++++++++----
 2 files changed, 33 insertions(+), 4 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
index afca5fb6f9..89916316dd 100644
--- a/zuul-core/src/main/java/com/netflix/zuul/context/SessionContext.java
+++ b/zuul-core/src/main/java/com/netflix/zuul/context/SessionContext.java
@@ -151,6 +151,23 @@ public <T> T getOrDefault(Key<T> key, T defaultValue) {
         return defaultValue;
     }
 
+    /**
+     * {@inheritDoc}
+     *
+     * <p>This method exists for static analysis.
+     */
+    @Override
+    public boolean containsKey(Object key) {
+        return super.containsKey(key);
+    }
+
+    /**
+     * Checks for the existence of the key in the context.
+     */
+    public <T> boolean containsKey(Key<T> key) {
+        return typedMap.containsKey(Objects.requireNonNull(key, "key"));
+    }
+
     /**
      * {@inheritDoc}
      *
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
index 3322744291..1ce2bbf932 100644
--- a/zuul-core/src/test/java/com/netflix/zuul/context/SessionContextTest.java
+++ b/zuul-core/src/test/java/com/netflix/zuul/context/SessionContextTest.java
@@ -15,16 +15,14 @@
  */
 package com.netflix.zuul.context;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
 import com.google.common.truth.Truth;
 import com.netflix.zuul.context.SessionContext.Key;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.junit.jupiter.MockitoExtension;
 
+import static org.junit.jupiter.api.Assertions.*;
+
 @ExtendWith(MockitoExtension.class)
 class SessionContextTest {
 
@@ -107,4 +105,18 @@ void remove() {
         Truth.assertThat(context.get(key)).isNull();
         Truth.assertThat(val).isEqualTo("bar");
     }
+
+    @Test
+    void containsKey() {
+        SessionContext context = new SessionContext();
+        Key<String> key = SessionContext.newKey("foo");
+        context.put(key, "bar");
+
+        Truth.assertThat(context.containsKey(key)).isTrue();
+
+        String val = context.remove(key);
+        Truth.assertThat(val).isEqualTo("bar");
+
+        Truth.assertThat(context.containsKey(key)).isFalse();
+    }
 }