Fixed test

DataDog · Jun 25, 2024 · 36af77f · 36af77f
1 parent b3305e9
commit 36af77f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/...agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy
@@ -24,7 +24,6 @@ import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter
 import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase
 import datadog.trace.test.util.DDSpecification
 
-import java.util.function.BiConsumer
 import java.util.function.BiFunction
 import java.util.function.Function
 import java.util.function.Supplier
@@ -80,8 +79,8 @@ class GatewayBridgeSpecification extends DDSpecification {
   BiFunction<RequestContext, String, Flow<Void>> grpcServerMethodCB
   BiFunction<RequestContext, Object, Flow<Void>> grpcServerRequestMessageCB
   BiFunction<RequestContext, Map<String, Object>, Flow<Void>> graphqlServerRequestMessageCB
-  BiConsumer<RequestContext, String> databaseConnectionCB
-  BiConsumer<RequestContext, String> databaseSqlQueryCB
+  BiFunction<RequestContext, String, Flow<Void>> databaseConnectionCB
+  BiFunction<RequestContext, String, Flow<Void>> databaseSqlQueryCB
 
   void setup() {
     callInitAndCaptureCBs()

diff --git a/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java b/internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java
@@ -201,9 +201,9 @@ public void testNormalCalls() {
     assertThat(cbp.getCallback(events.graphqlServerRequestMessage()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);
     ss.registerCallback(events.databaseConnection(), callback);
-    cbp.getCallback(events.databaseConnection()).accept(null, null);
+    cbp.getCallback(events.databaseConnection()).apply(null, null);
     ss.registerCallback(events.databaseSqlQuery(), callback);
-    cbp.getCallback(events.databaseSqlQuery()).accept(null, null);
+    cbp.getCallback(events.databaseSqlQuery()).apply(null, null);
     assertThat(callback.count).isEqualTo(Events.MAX_EVENTS);
   }
 
@@ -257,9 +257,9 @@ public void testThrowableBlocking() {
     assertThat(cbp.getCallback(events.graphqlServerRequestMessage()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);
     ss.registerCallback(events.databaseConnection(), throwback);
-    cbp.getCallback(events.databaseConnection()).accept(null, null);
+    cbp.getCallback(events.databaseConnection()).apply(null, null);
     ss.registerCallback(events.databaseSqlQuery(), throwback);
-    cbp.getCallback(events.databaseSqlQuery()).accept(null, null);
+    cbp.getCallback(events.databaseSqlQuery()).apply(null, null);
     assertThat(throwback.count).isEqualTo(Events.MAX_EVENTS);
   }