Fix instrumentation gateway tests

Author: manuel-alvarez-alvarez

internal-api/src/main/java/datadog/trace/api/gateway/InstrumentationGateway.java

@@ -22,6 +22,7 @@
 import static datadog.trace.api.gateway.Events.REQUEST_PATH_PARAMS_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_SESSION_ID;
 import static datadog.trace.api.gateway.Events.REQUEST_STARTED_ID;
+import static datadog.trace.api.gateway.Events.RESPONSE_BODY_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_DONE_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_HEADER_ID;
 import static datadog.trace.api.gateway.Events.RESPONSE_STARTED_ID;
@@ -346,6 +347,7 @@ public Flow<Void> apply(RequestContext ctx, StoredBodySupplier storedBodySupplie
       case GRPC_SERVER_REQUEST_MESSAGE_ID:
       case GRAPHQL_SERVER_REQUEST_MESSAGE_ID:
       case REQUEST_BODY_CONVERTED_ID:
+      case RESPONSE_BODY_ID:
         return (C)
             new BiFunction<RequestContext, Object, Flow<Void>>() {
               @Override

internal-api/src/test/java/datadog/trace/api/gateway/InstrumentationGatewayTest.java

@@ -190,6 +190,9 @@ public void testNormalCalls() {
     ss.registerCallback(events.requestBodyProcessed(), callback);
     assertThat(cbp.getCallback(events.requestBodyProcessed()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);
+    ss.registerCallback(events.responseBody(), callback);
+    assertThat(cbp.getCallback(events.responseBody()).apply(null, null).getAction())
+        .isEqualTo(Flow.Action.Noop.INSTANCE);
     ss.registerCallback(events.grpcServerMethod(), callback);
     assertThat(cbp.getCallback(events.grpcServerMethod()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);
@@ -260,6 +263,9 @@ public void testThrowableBlocking() {
     ss.registerCallback(events.requestBodyProcessed(), throwback);
     assertThat(cbp.getCallback(events.requestBodyProcessed()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);
+    ss.registerCallback(events.responseBody(), throwback);
+    assertThat(cbp.getCallback(events.responseBody()).apply(null, null).getAction())
+        .isEqualTo(Flow.Action.Noop.INSTANCE);
     ss.registerCallback(events.grpcServerMethod(), throwback);
     assertThat(cbp.getCallback(events.grpcServerMethod()).apply(null, null).getAction())
         .isEqualTo(Flow.Action.Noop.INSTANCE);