add more tests

Author: samuelcolvin

instrumentation/opentelemetry-instrumentation-flask/tests/base_test.py

@@ -88,6 +88,16 @@ def _custom_response_headers():
         resp.headers["my-secret-header"] = "my-secret-value"
         return resp
 
+    @staticmethod
+    def _repeat_custom_response_headers():
+        headers = {
+            "content-type": "text/plain; charset=utf-8",
+            "my-custom-header": [
+                "my-custom-value-1", "my-custom-header-2"
+            ],
+        }
+        return flask.Response("test response", headers=headers)
+
     def _common_initialization(self):
         def excluded_endpoint():
             return "excluded"
@@ -106,6 +116,9 @@ def excluded2_endpoint():
         self.app.route("/test_custom_response_headers")(
             self._custom_response_headers
         )
+        self.app.route("/test_repeat_custom_response_headers")(
+            self._repeat_custom_response_headers
+        )
 
         # pylint: disable=attribute-defined-outside-init
         self.client = Client(self.app, Response)

instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py

@@ -671,6 +671,22 @@ def test_custom_request_header_added_in_server_span(self):
         self.assertEqual(span.kind, trace.SpanKind.SERVER)
         self.assertSpanHasAttributes(span, expected)
 
+    def test_repeat_custom_request_header_added_in_server_span(self):
+        headers = [
+            ("Custom-Test-Header-1", "Test Value 1"),
+            ("Custom-Test-Header-1", "Test Value 2"),
+        ]
+        resp = self.client.get("/hello/123", headers=headers)
+        self.assertEqual(200, resp.status_code)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.request.header.custom_test_header_1": (
+                "Test Value 1, Test Value 2",
+            ),
+        }
+        self.assertEqual(span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+
     def test_custom_request_header_not_added_in_internal_span(self):
         tracer = trace.get_tracer(__name__)
         with tracer.start_as_current_span("test", kind=trace.SpanKind.SERVER):
@@ -724,6 +740,21 @@ def test_custom_response_header_added_in_server_span(self):
         self.assertEqual(span.kind, trace.SpanKind.SERVER)
         self.assertSpanHasAttributes(span, expected)
 
+    def test_repeat_custom_response_header_added_in_server_span(self):
+        resp = self.client.get("/test_repeat_custom_response_headers")
+        self.assertEqual(resp.status_code, 200)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.response.header.content_type": (
+                "text/plain; charset=utf-8",
+            ),
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-header-2",
+            ),
+        }
+        self.assertEqual(span.kind, trace.SpanKind.SERVER)
+        self.assertSpanHasAttributes(span, expected)
+
     def test_custom_response_header_not_added_in_internal_span(self):
         tracer = trace.get_tracer(__name__)
         with tracer.start_as_current_span("test", kind=trace.SpanKind.SERVER):

instrumentation/opentelemetry-instrumentation-wsgi/tests/test_wsgi_middleware.py

@@ -115,6 +115,18 @@ def wsgi_with_custom_response_headers(environ, start_response):
     return [b"*"]
 
 
+def wsgi_with_repeat_custom_response_headers(environ, start_response):
+    assert isinstance(environ, dict)
+    start_response(
+        "200 OK",
+        [
+            ("my-custom-header", "my-custom-value-1"),
+            ("my-custom-header", "my-custom-value-2"),
+        ],
+    )
+    return [b"*"]
+
+
 _expected_metric_names = [
     "http.server.active_requests",
     "http.server.duration",
@@ -711,6 +723,26 @@ def test_custom_response_headers_not_added_in_internal_span(self):
             for key, _ in not_expected.items():
                 self.assertNotIn(key, span.attributes)
 
+    @mock.patch.dict(
+        "os.environ",
+        {
+            OTEL_INSTRUMENTATION_HTTP_CAPTURE_HEADERS_SERVER_RESPONSE: "my-custom-header",
+        },
+    )
+    def test_custom_response_headers_added_in_server_span(self):
+        app = otel_wsgi.OpenTelemetryMiddleware(
+            wsgi_with_repeat_custom_response_headers
+        )
+        response = app(self.environ, self.start_response)
+        self.iterate_response(response)
+        span = self.memory_exporter.get_finished_spans()[0]
+        expected = {
+            "http.response.header.my_custom_header": (
+                "my-custom-value-1,my-custom-value-2",
+            ),
+        }
+        self.assertSpanHasAttributes(span, expected)
+
 
 if __name__ == "__main__":
     unittest.main()