add test case

oliver-zhang · Jul 5, 2023 · 9d82f76 · 9d82f76
1 parent 41c461d
commit 9d82f76
Show file tree

Hide file tree

Showing 4 changed files with 130 additions and 0 deletions.
diff --git a/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/AbstractServlet3Test.groovy b/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/AbstractServlet3Test.groovy
@@ -66,6 +66,32 @@ abstract class AbstractServlet3Test<SERVER, CONTEXT> extends HttpServerTest<SERV
         + "<p>test works</p>\n"
         + "</body>\n"
         + "</html>")
+  public static final ServerEndpoint HTML_PRINT_WRITER_WITH_OTHER_HEAD =
+    new ServerEndpoint("HTML_PRINT_WRITER_WITH_OTHER_HEAD",  "htmlPrintWriterWithOtherHead",
+      200,
+      "<!DOCTYPE html>\n"
+        + "<html lang=\"en\">\n"
+        + "<head lang=\"en\">\n"
+        + "  <meta charset=\"UTF-8\">\n"
+        + "  <title>Title</title>\n"
+        + "</head>\n"
+        + "<body>\n"
+        + "<p>test works</p>\n"
+        + "</body>\n"
+        + "</html>")
+  public static final ServerEndpoint HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD =
+    new ServerEndpoint("HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD",  "htmlServletOutputStreamWithOtherHead",
+      200,
+      "<!DOCTYPE html>\n"
+        + "<html lang=\"en\">\n"
+        + "<head lang=\"en\">\n"
+        + "  <meta charset=\"UTF-8\">\n"
+        + "  <title>Title</title>\n"
+        + "</head>\n"
+        + "<body>\n"
+        + "<p>test works</p>\n"
+        + "</body>\n"
+        + "</html>")
   protected void setupServlets(CONTEXT context) {
     def servlet = servlet()
 
@@ -80,6 +106,8 @@ abstract class AbstractServlet3Test<SERVER, CONTEXT> extends HttpServerTest<SERV
     addServlet(context, CAPTURE_PARAMETERS.path, servlet)
     addServlet(context, HTML_PRINT_WRITER.path, servlet)
     addServlet(context, HTML_SERVLET_OUTPUT_STREAM.path, servlet)
+    addServlet(context, HTML_PRINT_WRITER_WITH_OTHER_HEAD.path,servlet)
+    addServlet(context, HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD.path,servlet)
   }
 
   protected ServerEndpoint lastRequest
@@ -215,4 +243,89 @@ abstract class AbstractServlet3Test<SERVER, CONTEXT> extends HttpServerTest<SERV
       }
     }
   }
+
+  def "snippet injection with PrintWriterWithOtherHead"() {
+    setup:
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test </script>")
+    def request = request(HTML_PRINT_WRITER_WITH_OTHER_HEAD, "GET")
+    def response = client.execute(request).aggregate().join()
+
+    expect:
+    response.status().code() == HTML_PRINT_WRITER_WITH_OTHER_HEAD.status
+    String result = "<!DOCTYPE html>\n" +
+      "<html lang=\"en\">\n" +
+      "<head lang=\"en\">\n" +
+      "  <script type=\"text/javascript\"> Test </script>\n" +
+      "  <meta charset=\"UTF-8\">\n" +
+      "  <title>Title</title>\n" +
+      "</head>\n" +
+      "<body>\n" +
+      "<p>test works</p>\n" +
+      "</body>\n" +
+      "</html>"
+    println "result from call ${response.contentUtf8()}"
+    response.contentUtf8() == result
+    response.headers().contentLength() == result.length()
+
+    cleanup:
+    ExperimentalSnippetHolder.setSnippet("")
+
+    def expectedRoute = expectedHttpRoute(HTML_PRINT_WRITER_WITH_OTHER_HEAD)
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          name "GET" + (expectedRoute != null ? " " + expectedRoute : "")
+          kind SpanKind.SERVER
+          hasNoParent()
+        }
+        span(1) {
+          name "controller"
+          kind SpanKind.INTERNAL
+          childOf span(0)
+        }
+      }
+    }
+  }
+
+  def "snippet injection with ServletOutputStreamWithOtherHead"() {
+    setup:
+    ExperimentalSnippetHolder.setSnippet("\n  <script type=\"text/javascript\"> Test Test</script>")
+    def request = request(HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD, "GET")
+    def response = client.execute(request).aggregate().join()
+
+    expect:
+    response.status().code() == HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD.status
+    String result = "<!DOCTYPE html>\n" +
+      "<html lang=\"en\">\n" +
+      "<head lang=\"en\">\n" +
+      "  <script type=\"text/javascript\"> Test Test</script>\n" +
+      "  <meta charset=\"UTF-8\">\n" +
+      "  <title>Title</title>\n" +
+      "</head>\n" +
+      "<body>\n" +
+      "<p>test works</p>\n" +
+      "</body>\n" +
+      "</html>"
+    response.contentUtf8() == result
+    response.headers().contentLength() == result.length()
+
+    cleanup:
+    ExperimentalSnippetHolder.setSnippet("")
+
+    def expectedRoute = expectedHttpRoute(HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD)
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          name "GET" + (expectedRoute != null ? " " + expectedRoute : "")
+          kind SpanKind.SERVER
+          hasNoParent()
+        }
+        span(1) {
+          name "controller"
+          kind SpanKind.INTERNAL
+          childOf span(0)
+        }
+      }
+    }
+  }
 }
diff --git a/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/JettyServlet3Test.groovy b/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/JettyServlet3Test.groovy
@@ -266,6 +266,8 @@ class JettyServlet3TestDispatchAsync extends JettyDispatchTest {
   protected void setupServlets(ServletContextHandler context) {
     super.setupServlets(context)
     addServlet(context, "/dispatch" + HTML_PRINT_WRITER.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + HTML_PRINT_WRITER_WITH_OTHER_HEAD.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD.path, TestServlet3.DispatchAsync)
     addServlet(context, "/dispatch" + HTML_SERVLET_OUTPUT_STREAM.path, TestServlet3.DispatchAsync)
     addServlet(context, "/dispatch" + SUCCESS.path, TestServlet3.DispatchAsync)
     addServlet(context, "/dispatch" + QUERY_PARAM.path, TestServlet3.DispatchAsync)

diff --git a/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/TestServlet3.groovy b/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/TestServlet3.groovy
@@ -84,6 +84,19 @@ class TestServlet3 {
             byte[] body = endpoint.body.getBytes()
             resp.getOutputStream().write(body, 0, body.length)
             break
+          case AbstractServlet3Test.HTML_PRINT_WRITER_WITH_OTHER_HEAD:
+            resp.contentType = "text/html"
+            resp.status = endpoint.status
+            resp.setContentLengthLong(endpoint.body.length())
+            resp.writer.print(endpoint.body)
+            break
+          case AbstractServlet3Test.HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD:
+            resp.contentType = "text/html"
+            resp.status = endpoint.status
+            resp.setContentLength(endpoint.body.length())
+            byte[] body = endpoint.body.getBytes()
+            resp.getOutputStream().write(body, 0, body.length)
+            break
         }
       }
     }

diff --git a/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/TomcatServlet3Test.groovy b/instrumentation/servlet/servlet-3.0/javaagent/src/test/groovy/TomcatServlet3Test.groovy
@@ -363,6 +363,8 @@ class TomcatServlet3TestForward extends TomcatDispatchTest {
     addServlet(context, "/dispatch" + INDEXED_CHILD.path, RequestDispatcherServlet.Forward)
     addServlet(context, "/dispatch" + HTML_PRINT_WRITER.path, RequestDispatcherServlet.Forward)
     addServlet(context, "/dispatch" + HTML_SERVLET_OUTPUT_STREAM.path, RequestDispatcherServlet.Forward)
+    addServlet(context, "/dispatch" + HTML_SERVLET_OUTPUT_STREAM_WITH_OTHER_HEAD.path, RequestDispatcherServlet.Forward)
+    addServlet(context, "/dispatch" + HTML_PRINT_WRITER_WITH_OTHER_HEAD.path, RequestDispatcherServlet.Forward)
   }
 }