Fix failing test

DataDog · Jul 10, 2024 · c399041 · c399041
1 parent 03922de
commit c399041
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp-2/src/test/groovy/IastOkHttp2InstrumentationTest.groovy
@@ -2,6 +2,7 @@ import com.squareup.okhttp.OkHttpClient
 import com.squareup.okhttp.Request
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.iast.InstrumentationBridge
+import datadog.trace.api.iast.propagation.CodecModule
 import datadog.trace.api.iast.propagation.PropagationModule
 import datadog.trace.api.iast.sink.SsrfModule
 import datadog.trace.instrumentation.okhttp2.IastHttpUrlInstrumentation
@@ -61,18 +62,20 @@ class IastOkHttp2InstrumentationTest extends AgentTestRunner {
   }
 
   private void mockPropagation() {
-    final propagation = Mock(PropagationModule) {
-      taintObjectIfAnyTainted(_, _) >> {
-        if ((it[1] as List).any { input -> tainteds.containsKey(input) }) {
-          tainteds.put(it[0], null)
-        }
-      }
-      taintObjectIfTainted(_, _) >> {
-        if (tainteds.containsKey(it[1])) {
-          tainteds.put(it[0], null)
-        }
+    final Closure taint = { target, inputs ->
+      if (inputs.any { input -> tainteds.containsKey(input) }) {
+        tainteds.put(target, null)
       }
     }
+
+    final propagation = Mock(PropagationModule) {
+      taintStringIfTainted(_, _) >> { taint(it[0], [it[1]]) }
+    }
     InstrumentationBridge.registerIastModule(propagation)
+
+    final codec = Mock(CodecModule) {
+      onUrlCreate(_, _) >> { taint(it[0], it[1] as List) }
+    }
+    InstrumentationBridge.registerIastModule(codec)
   }
 }