diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java b/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java
index 1408afa10c7cf..06d223907082e 100644
--- a/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java
+++ b/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java
@@ -26,13 +26,12 @@ public class QueryGroupThreadContextStatePropagator implements ThreadContextStat
      * @return the map of header and their values to be propagated across request threadContexts
      */
     @Override
+    @SuppressWarnings("removal")
     public Map<String, Object> transients(Map<String, Object> source) {
         final Map<String, Object> transientHeaders = new HashMap<>();
 
         for (String headerName : PROPAGATED_HEADERS) {
-            if (source.containsKey(headerName)) {
-                transientHeaders.put(headerName, source.get(headerName));
-            }
+            transientHeaders.compute(headerName, (k, v) -> source.get(headerName));
         }
         return transientHeaders;
     }
@@ -42,13 +41,12 @@ public Map<String, Object> transients(Map<String, Object> source) {
      * @return map of header and their values to be propagated across nodes
      */
     @Override
+    @SuppressWarnings("removal")
     public Map<String, String> headers(Map<String, Object> source) {
         final Map<String, String> propagatedHeaders = new HashMap<>();
 
         for (String headerName : PROPAGATED_HEADERS) {
-            if (source.containsKey(headerName)) {
-                propagatedHeaders.put(headerName, (String) source.get(headerName));
-            }
+            propagatedHeaders.compute(headerName, (k, v) -> (String) source.get(headerName));
         }
         return propagatedHeaders;
     }
diff --git a/server/src/main/java/org/opensearch/wlm/package-info.java b/server/src/main/java/org/opensearch/wlm/package-info.java
index f6a8008b6374e..fa4731d95cc34 100644
--- a/server/src/main/java/org/opensearch/wlm/package-info.java
+++ b/server/src/main/java/org/opensearch/wlm/package-info.java
@@ -6,8 +6,8 @@
  * compatible open source license.
  */
 
-package org.opensearch.wlm;
-
 /**
  * This package contains workload management constructs
  */
+
+package org.opensearch.wlm;
diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTest.java b/server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTests.java
similarity index 92%
rename from server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTest.java
rename to server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTests.java
index 60c8f2605914f..ad5d7f569a56e 100644
--- a/server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTest.java
+++ b/server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTests.java
@@ -12,7 +12,7 @@
 
 import java.util.Map;
 
-public class QueryGroupThreadContextStatePropagatorTest extends OpenSearchTestCase {
+public class QueryGroupThreadContextStatePropagatorTests extends OpenSearchTestCase {
 
     public void testTransients() {
         QueryGroupThreadContextStatePropagator sut = new QueryGroupThreadContextStatePropagator();