feat(messaging): add toggle for command request header propagation (#…

…1114) - Introduce ENABLED_KEY constant to control propagation feature - Add conditional check to enable/disable header injection - Update tests to cover both enabled and disabled scenarios
Ahoo-Wang · Jan 9, 2025 · 000072e · 000072e
1 parent 036ef0f
commit 000072e
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/wow-core/src/main/kotlin/me/ahoo/wow/messaging/propagation/CommandRequestHeaderPropagator.kt b/wow-core/src/main/kotlin/me/ahoo/wow/messaging/propagation/CommandRequestHeaderPropagator.kt
@@ -18,6 +18,7 @@ import me.ahoo.wow.api.messaging.Message
 
 class CommandRequestHeaderPropagator : MessagePropagator {
     companion object {
+        const val ENABLED_KEY = "wow.messaging.propagation.request"
         private const val USER_AGENT = "user_agent"
         private const val REMOTE_IP = "remote_ip"
         val Header.userAgent: String?
@@ -39,7 +40,12 @@ class CommandRequestHeaderPropagator : MessagePropagator {
         }
     }
 
+    private val enabled: Boolean = System.getProperty(ENABLED_KEY)?.toBoolean() != false
+
     override fun inject(header: Header, upstream: Message<*, *>) {
+        if (!enabled) {
+            return
+        }
         upstream.header.userAgent?.let {
             header.withUserAgent(it)
         }

diff --git a/...e/src/test/kotlin/me/ahoo/wow/messaging/propagation/CommandRequestHeaderPropagatorTest.kt b/...e/src/test/kotlin/me/ahoo/wow/messaging/propagation/CommandRequestHeaderPropagatorTest.kt
@@ -48,4 +48,32 @@ class CommandRequestHeaderPropagatorTest {
         assertThat(injectedHeader.userAgent, equalTo(null))
         assertThat(injectedHeader.remoteIp, equalTo(null))
     }
+
+    @Test
+    fun injectDisabled() {
+        System.setProperty(CommandRequestHeaderPropagator.ENABLED_KEY, "false")
+        val injectedHeader = DefaultHeader.empty()
+        val upstreamMessage =
+            MockCreateAggregate(GlobalIdGenerator.generateAsString(), GlobalIdGenerator.generateAsString())
+                .toCommandMessage()
+        upstreamMessage.header.withUserAgent("userAgent").withRemoteIp("remoteIp")
+        CommandRequestHeaderPropagator().inject(injectedHeader, upstreamMessage)
+        assertThat(injectedHeader.userAgent, equalTo(null))
+        assertThat(injectedHeader.remoteIp, equalTo(null))
+        System.clearProperty(CommandRequestHeaderPropagator.ENABLED_KEY)
+    }
+
+    @Test
+    fun injectEnabled() {
+        System.setProperty(CommandRequestHeaderPropagator.ENABLED_KEY, "true")
+        val injectedHeader = DefaultHeader.empty()
+        val upstreamMessage =
+            MockCreateAggregate(GlobalIdGenerator.generateAsString(), GlobalIdGenerator.generateAsString())
+                .toCommandMessage()
+        upstreamMessage.header.withUserAgent("userAgent").withRemoteIp("remoteIp")
+        CommandRequestHeaderPropagator().inject(injectedHeader, upstreamMessage)
+        assertThat(injectedHeader.userAgent, equalTo(upstreamMessage.header.userAgent))
+        assertThat(injectedHeader.remoteIp, equalTo(upstreamMessage.header.remoteIp))
+        System.clearProperty(CommandRequestHeaderPropagator.ENABLED_KEY)
+    }
 }