Fix more tests

DataDog · Jul 15, 2024 · b3dfd5a · b3dfd5a
1 parent a2f79a9
commit b3dfd5a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 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
@@ -69,13 +69,13 @@ class IastOkHttp2InstrumentationTest extends AgentTestRunner {
     }
 
     final propagation = Mock(PropagationModule) {
-      taintStringIfTainted(_, _) >> { taint(it[0], [it[1]]) }
-      taintObjectIfTainted(_, _) >> { taint(it[0], [it[1]]) }
+      taintStringIfTainted(*_) >> { taint(it[0], [it[1]]) }
+      taintObjectIfTainted(*_) >> { taint(it[0], [it[1]]) }
     }
     InstrumentationBridge.registerIastModule(propagation)
 
     final codec = Mock(CodecModule) {
-      onUrlCreate(_, _) >> { taint(it[0], it[1] as List) }
+      onUrlCreate(*_) >> { taint(it[0], it[1] as List) }
     }
     InstrumentationBridge.registerIastModule(codec)
   }

diff --git a/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy b/dd-java-agent/instrumentation/okhttp-3/src/test/groovy/IastOkHttp3InstrumentationTest.groovy
@@ -1,5 +1,6 @@
 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.okhttp3.IastHttpUrlInstrumentation
@@ -66,15 +67,21 @@ class IastOkHttp3InstrumentationTest extends AgentTestRunner {
   }
 
   private void mockPropagation() {
-    final Closure taint = { target, input ->
-      if (tainteds.containsKey(input)) {
+    final Closure taint = { target, inputs ->
+      if (inputs.any { input -> tainteds.containsKey(input) }) {
         tainteds.put(target, null)
       }
     }
+
     final propagation = Mock(PropagationModule) {
-      taintObjectIfTainted(_, _) >> { taint(it[0], it[1]) }
-      taintStringIfTainted(_, _) >> { taint(it[0], it[1]) }
+      taintStringIfTainted(*_) >> { taint(it[0], [it[1]]) }
+      taintObjectIfTainted(*_) >> { taint(it[0], [it[1]]) }
     }
     InstrumentationBridge.registerIastModule(propagation)
+
+    final codec = Mock(CodecModule) {
+      onUrlCreate(*_) >> { taint(it[0], it[1] as List) }
+    }
+    InstrumentationBridge.registerIastModule(codec)
   }
 }