Refactor code

Signed-off-by: Gagan Juneja <[email protected]>
opensearch-project · Jan 8, 2024 · 9b9a895 · 9b9a895
1 parent 220aff1
commit 9b9a895
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 43 deletions.
diff --git a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContextStatePropagator.java b/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContextStatePropagator.java
@@ -26,6 +26,7 @@ public interface ThreadContextStatePropagator {
      * @param source          current context transient headers
      * @return the list of transient headers that needs to be propagated from current context to new thread context
      */
+    @Deprecated(since = "2.12.0", forRemoval = true)
     Map<String, Object> transients(Map<String, Object> source);
 
     /**
@@ -45,6 +46,7 @@ default Map<String, Object> transients(Map<String, Object> source, boolean isSys
      * @param source          current context headers
      * @return the list of request headers that needs to be propagated from current context to request
      */
+    @Deprecated(since = "2.12.0", forRemoval = true)
     Map<String, String> headers(Map<String, Object> source);
 
     /**

diff --git a/...rc/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java b/...rc/main/java/org/opensearch/telemetry/tracing/ThreadContextBasedTracerContextStorage.java
@@ -85,6 +85,11 @@ public Map<String, String> headers(Map<String, Object> source) {
         return headers;
     }
 
+    @Override
+    public Map<String, String> headers(Map<String, Object> source, boolean isSystemContext) {
+        return headers(source);
+    }
+
     Span getCurrentSpan(String key) {
         SpanReference currentSpanRef = threadContext.getTransient(key);
         return (currentSpanRef == null) ? null : currentSpanRef.getSpan();

diff --git a/server/src/test/java/org/opensearch/common/util/concurrent/ThreadContextTests.java b/server/src/test/java/org/opensearch/common/util/concurrent/ThreadContextTests.java
@@ -44,6 +44,8 @@
 import java.util.Map;
 import java.util.function.Supplier;
 
+import org.mockito.Mockito;
+
 import static org.opensearch.tasks.TaskResourceTrackingService.TASK_ID;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
@@ -742,8 +744,18 @@ public void testMarkAsSystemContext() throws IOException {
 
     public void testSystemContextWithPropagator() {
         Settings build = Settings.builder().put("request.headers.default", "1").build();
+        Map<String, Object> transientHeaderMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, Object> transientHeaderTransformedMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, Object> headerMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, String> headerTransformedMap = Collections.singletonMap("test_transient_propagation_key", "test");
         ThreadContext threadContext = new ThreadContext(build);
-        threadContext.registerThreadContextStatePropagator(createDummyPropagator(("test_transient_propagation_key")));
+        ThreadContextStatePropagator mockPropagator = Mockito.mock(ThreadContextStatePropagator.class);
+        Mockito.when(mockPropagator.transients(transientHeaderMap, true)).thenReturn(Collections.emptyMap());
+        Mockito.when(mockPropagator.transients(transientHeaderMap, false)).thenReturn(transientHeaderTransformedMap);
+
+        Mockito.when(mockPropagator.headers(headerMap, true)).thenReturn(headerTransformedMap);
+        Mockito.when(mockPropagator.headers(headerMap, false)).thenReturn(headerTransformedMap);
+        threadContext.registerThreadContextStatePropagator(mockPropagator);
         threadContext.putHeader("foo", "bar");
         threadContext.putTransient("test_transient_propagation_key", 1);
         assertEquals(Integer.valueOf(1), threadContext.getTransient("test_transient_propagation_key"));
@@ -762,8 +774,18 @@ public void testSystemContextWithPropagator() {
 
     public void testSerializeSystemContext() throws IOException {
         Settings build = Settings.builder().put("request.headers.default", "1").build();
+        Map<String, Object> transientHeaderMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, Object> transientHeaderTransformedMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, Object> headerMap = Collections.singletonMap("test_transient_propagation_key", "test");
+        Map<String, String> headerTransformedMap = Collections.singletonMap("test_transient_propagation_key", "test");
         ThreadContext threadContext = new ThreadContext(build);
-        threadContext.registerThreadContextStatePropagator(createDummyPropagator(("test_transient_propagation_key")));
+        ThreadContextStatePropagator mockPropagator = Mockito.mock(ThreadContextStatePropagator.class);
+        Mockito.when(mockPropagator.transients(transientHeaderMap, true)).thenReturn(Collections.emptyMap());
+        Mockito.when(mockPropagator.transients(transientHeaderMap, false)).thenReturn(transientHeaderTransformedMap);
+
+        Mockito.when(mockPropagator.headers(headerMap, true)).thenReturn(headerTransformedMap);
+        Mockito.when(mockPropagator.headers(headerMap, false)).thenReturn(headerTransformedMap);
+        threadContext.registerThreadContextStatePropagator(mockPropagator);
         threadContext.putHeader("foo", "bar");
         threadContext.putTransient("test_transient_propagation_key", "test");
         BytesStreamOutput out = new BytesStreamOutput();
@@ -785,47 +807,6 @@ public void testSerializeSystemContext() throws IOException {
         assertEquals("1", threadContext.getHeader("default"));
     }
 
-    private ThreadContextStatePropagator createDummyPropagator(final String key) {
-        return new ThreadContextStatePropagator() {
-
-            @Override
-            public Map<String, Object> transients(Map<String, Object> source) {
-                Map<String, Object> transients = new HashMap<>();
-                if (source.containsKey(key)) {
-                    transients.put(key, source.get(key));
-                }
-                return transients;
-            }
-
-            @Override
-            public Map<String, Object> transients(Map<String, Object> source, boolean isSystemContext) {
-                if (isSystemContext == true) {
-                    return Collections.emptyMap();
-                } else {
-                    return transients(source);
-                }
-            }
-
-            @Override
-            public Map<String, String> headers(Map<String, Object> source) {
-                Map<String, String> headers = new HashMap<>();
-                if (source.containsKey(key)) {
-                    headers.put(key, (String) source.get(key));
-                }
-                return headers;
-            }
-
-            @Override
-            public Map<String, String> headers(Map<String, Object> source, boolean isSystemContext) {
-                if (isSystemContext == true) {
-                    return Collections.emptyMap();
-                } else {
-                    return headers(source);
-                }
-            }
-        };
-    }
-
     public void testPutHeaders() {
         Settings build = Settings.builder().put("request.headers.default", "1").build();
         ThreadContext threadContext = new ThreadContext(build);