Chunked encoding didnt work as expected

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy

@@ -51,14 +51,6 @@ abstract class HttpClientTest<DECORATOR extends HttpClientDecorator> extends Age
         handleDistributedRequest()
         redirect(server.address.resolve("/circular-redirect").toURL().toString())
       }
-      prefix("chunked-response") {
-        handleDistributedRequest()
-        response.status(200)
-          .startChunked()
-          .sendChunk("first part")
-          .sendChunk("-second part")
-          .finishChunks()
-      }
     }
   }
 
@@ -147,26 +139,6 @@ abstract class HttpClientTest<DECORATOR extends HttpClientDecorator> extends Age
     method = "HEAD"
   }
 
-  def "handle chunked response"() {
-    when:
-    def status = runUnderTrace("parent") {
-      doRequest(method, server.address.resolve("/chunked-response"))
-    }
-
-    then:
-    status == 200
-    assertTraces(2) {
-      server.distributedRequestTrace(it, 0, trace(1).last())
-      trace(1, size(2)) {
-        basicSpan(it, 0, "parent")
-        clientSpan(it, 1, span(0), method, false, false, server.address.resolve("/chunked-response"))
-      }
-    }
-
-    where:
-    method << BODY_METHODS
-  }
-
   def "trace request without propagation"() {
     when:
     def status = withConfigOverride(Config.HTTP_CLIENT_HOST_SPLIT_BY_DOMAIN, "$renameService") {

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/server/http/TestHttpServer.groovy

@@ -302,29 +302,6 @@ class TestHttpServer implements AutoCloseable {
         resp.setContentLength(body.bytes.length)
         resp.writer.print(body)
       }
-
-      ResponseApi startChunked() {
-        assert !req.handled
-        resp.setHeader("Transfer-Encoding", "chunked")
-        resp.status = status
-        return this
-      }
-
-      ResponseApi sendChunk(String partial) {
-        assert !req.handled
-        resp.writer.print(Integer.toHexString(partial.length()).toUpperCase())
-        resp.writer.print("\r\n")
-        resp.writer.print(partial)
-        resp.writer.print("\r\n")
-        resp.writer.flush()
-        return this
-      }
-
-      void finishChunks() {
-        resp.writer.print("0\r\n\r\n")
-        resp.writer.flush()
-        req.handled = true
-      }
     }
 
     static class Headers {

dd-java-agent/testing/src/test/groovy/server/ServerTest.groovy

@@ -1,14 +1,14 @@
 package server
 
-
 import datadog.trace.agent.test.AgentTestRunner
-import datadog.trace.agent.test.asserts.ListWriterAssert
 import datadog.trace.agent.test.utils.OkHttpUtils
 import okhttp3.MultipartBody
 import okhttp3.Request
 import spock.lang.Shared
 
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
+import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
 /* Don't actually need AgentTestRunner, but it messes up the classloader for AgentTestRunnerTest if this runs first. */
 
@@ -312,7 +312,7 @@ class ServerTest extends AgentTestRunner {
     response.code() == 200
     response.body().string().trim() == "done"
 
-    ListWriterAssert.assertTraces(TEST_WRITER, 1) {
+    assertTraces(1) {
       server.distributedRequestTrace(it, 0)
     }
 
@@ -320,6 +320,129 @@ class ServerTest extends AgentTestRunner {
     server.stop()
   }
 
+  def "server ignores distributed request when header set"() {
+    setup:
+    def server = httpServer {
+      handlers {
+        all {
+          handleDistributedRequest()
+          response.send("done")
+        }
+      }
+    }
+
+    when:
+    def request = new Request.Builder()
+      .url("$server.address")
+      .header("is-dd-server", "false")
+      .get()
+      .build()
+
+    def response = runUnderTrace("parent") {
+      client.newCall(request).execute()
+    }
+
+    then:
+    response.code() == 200
+    response.body().string().trim() == "done"
+
+    assertTraces(1) {
+      trace(0, 1) {
+        basicSpan(it, 0, "parent")
+      }
+    }
+
+    cleanup:
+    server.stop()
+  }
+
+  def "server handles distributed request when header set"() {
+    setup:
+    def server = httpServer {
+      handlers {
+        all {
+          handleDistributedRequest()
+          response.send("done")
+        }
+      }
+    }
+
+    when:
+    def request = new Request.Builder()
+      .url("$server.address")
+      .header("is-dd-server", "true")
+      .get()
+      .build()
+
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == 200
+    response.body().string().trim() == "done"
+
+    // parent<->child relation can't be tested because okhttp isnt traced here
+    assertTraces(1) {
+      server.distributedRequestTrace(it, 0)
+    }
+
+    cleanup:
+    server.stop()
+  }
+
+
+  def "calling send() twice is an error"() {
+    setup:
+    def server = httpServer {
+      handlers {
+        all {
+          response.send()
+          response.send()
+        }
+      }
+    }
+
+    when:
+    def request = new Request.Builder()
+      .url("$server.address")
+      .get()
+      .build()
+
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == 500
+    response.message().startsWith("Server Error")
+
+    cleanup:
+    server.stop()
+  }
+
+  def "calling send() with null is an error"() {
+    setup:
+    def server = httpServer {
+      handlers {
+        all {
+          response.send(null)
+        }
+      }
+    }
+
+    when:
+    def request = new Request.Builder()
+      .url("$server.address")
+      .get()
+      .build()
+
+    def response = client.newCall(request).execute()
+
+    then:
+    response.code() == 500
+    response.message().startsWith("Server Error")
+
+    cleanup:
+    server.stop()
+  }
+
   def body() {
     return new MultipartBody.Builder().addFormDataPart("key", "value").build()
   }