websocket tests

Author: majorgreys

ext/opentelemetry-ext-asgi/src/opentelemetry/ext/asgi/__init__.py

@@ -28,8 +28,6 @@
 from opentelemetry.ext.asgi.version import __version__  # noqa
 from opentelemetry.trace.status import Status, StatusCanonicalCode
 
-_HTTP_VERSION_PREFIX = "HTTP/"
-
 
 def get_header_from_scope(scope: dict, header_name: str) -> typing.List[str]:
     """Retrieve a HTTP header value from the ASGI scope.
@@ -88,14 +86,16 @@ def collect_request_attributes(scope):
 
     result = {
         "component": scope.get("type"),
-        "http.method": scope.get("method"),
         "http.scheme": scope.get("scheme"),
         "http.host": server_host,
         "host.port": port,
         "http.flavor": scope.get("http_version"),
         "http.target": scope.get("path"),
         "http.url": http_url,
     }
+    http_method = scope.get("method")
+    if http_method:
+        result["http.method"] = http_method
     http_host_value = ",".join(get_header_from_scope(scope, "host"))
     if http_host_value:
         result["http.server_name"] = http_host_value
@@ -127,15 +127,10 @@ def set_status_code(span, status_code):
 
 
 def get_default_span_name(scope):
-    """Default implementation for name_callback, returns HTTP {METHOD_NAME or PATH}."""
-    span_name = "HTTP"
-
+    """Default implementation for name_callback"""
     method_or_path = scope.get("method") or scope.get("path")
 
-    if method_or_path:
-        return span_name + " " + method_or_path
-
-    return span_name
+    return method_or_path
 
 
 class OpenTelemetryMiddleware:
@@ -174,38 +169,32 @@ async def __call__(self, scope, receive, send):
 
         try:
             with self.tracer.start_as_current_span(
-                span_name + " (asgi.connection)",
+                span_name + " asgi",
                 kind=trace.SpanKind.SERVER,
                 attributes=collect_request_attributes(scope),
             ):
 
                 @wraps(receive)
                 async def wrapped_receive():
                     with self.tracer.start_as_current_span(
-                        span_name + " (asgi." + scope["type"] + ".receive)"
+                        span_name + " asgi." + scope["type"] + ".receive"
                     ) as receive_span:
                         payload = await receive()
                         if payload["type"] == "websocket.receive":
                             set_status_code(receive_span, 200)
-                            receive_span.set_attribute(
-                                "http.status_text", payload["text"]
-                            )
                         receive_span.set_attribute("type", payload["type"])
                     return payload
 
                 @wraps(send)
                 async def wrapped_send(payload):
                     with self.tracer.start_as_current_span(
-                        span_name + " (asgi." + scope["type"] + ".send)"
+                        span_name + " asgi." + scope["type"] + ".send"
                     ) as send_span:
                         if payload["type"] == "http.response.start":
                             status_code = payload["status"]
                             set_status_code(send_span, status_code)
                         elif payload["type"] == "websocket.send":
                             set_status_code(send_span, 200)
-                            send_span.set_attribute(
-                                "http.status_text", payload["text"]
-                            )
                         send_span.set_attribute("type", payload["type"])
                         await send(payload)
 

ext/opentelemetry-ext-asgi/tests/test_asgi_middleware.py

@@ -24,9 +24,7 @@
 )
 
 
-async def simple_asgi(scope, receive, send):
-    assert isinstance(scope, dict)
-    assert scope.get("type") == "http"
+async def http_app(scope, receive, send):
     payload = await receive()
     if payload.get("type") == "http.request":
         await send(
@@ -39,6 +37,28 @@ async def simple_asgi(scope, receive, send):
         await send({"type": "http.response.body", "body": b"*"})
 
 
+async def websocket_app(scope, receive, send):
+    while True:
+        payload = await receive()
+        if payload.get("type") == "websocket.connect":
+            await send({"type": "websocket.accept"})
+
+        if payload.get("type") == "websocket.receive":
+            if payload.get("text") == "ping":
+                await send({"type": "websocket.send", "text": "pong"})
+
+        if payload.get("type") == "websocket.disconnect":
+            break
+
+
+async def simple_asgi(scope, receive, send):
+    assert isinstance(scope, dict)
+    if scope["type"] == "http":
+        await http_app(scope, receive, send)
+    elif scope["type"] == "websocket":
+        await websocket_app(scope, receive, send)
+
+
 async def error_asgi(scope, receive, send):
     assert isinstance(scope, dict)
     assert scope.get("type") == "http"
@@ -91,25 +111,25 @@ def validate_outputs(self, outputs, error=None, modifiers=None):
         self.assertEqual(len(span_list), 4)
         expected = [
             {
-                "name": "HTTP GET (asgi.http.receive)",
+                "name": "GET asgi.http.receive",
                 "kind": trace_api.SpanKind.INTERNAL,
                 "attributes": {"type": "http.request"},
             },
             {
-                "name": "HTTP GET (asgi.http.send)",
+                "name": "GET asgi.http.send",
                 "kind": trace_api.SpanKind.INTERNAL,
                 "attributes": {
                     "http.status_code": 200,
                     "type": "http.response.start",
                 },
             },
             {
-                "name": "HTTP GET (asgi.http.send)",
+                "name": "GET asgi.http.send",
                 "kind": trace_api.SpanKind.INTERNAL,
                 "attributes": {"type": "http.response.body"},
             },
             {
-                "name": "HTTP GET (asgi.connection)",
+                "name": "GET asgi",
                 "kind": trace_api.SpanKind.SERVER,
                 "attributes": {
                     "component": "http",
@@ -228,41 +248,38 @@ def update_expected_user_agent(expected):
         self.validate_outputs(outputs, modifiers=[update_expected_user_agent])
 
     def test_websocket(self):
-        async def app_asgi(scope, receive, send):
-            assert scope["type"] == "websocket"
-            payload = await receive()
-            if payload.get("type") == "http.request":
-                await send(
-                    {
-                        "type": "http.response.start",
-                        "status": 200,
-                        "headers": [[b"Content-Type", b"text/plain"]],
-                    }
-                )
-            await send({"type": "http.response.body", "body": b"*"})
-
-        self.scope["type"] = "websocket"
-        app = otel_asgi.OpenTelemetryMiddleware(app_asgi)
+        self.scope = {
+            "type": "websocket",
+            "http_version": "1.1",
+            "scheme": "ws",
+            "path": "/",
+            "query_string": b"",
+            "headers": [],
+            "client": ("127.0.0.1", 32767),
+            "server": ("127.0.0.1", 80),
+        }
+        app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
         self.seed_app(app)
-        self.send_default_request()
+        self.send_input({"type": "websocket.connect"})
+        self.send_input({"type": "websocket.receive", "text": "ping"})
+        self.send_input({"type": "websocket.disconnect"})
+        outputs = self.get_all_output()
         span_list = self.memory_exporter.get_finished_spans()
-        self.assertEqual(len(span_list), 4)
+        self.assertEqual(len(span_list), 6)
         expected = [
-            "HTTP GET (asgi.websocket.receive)",
-            "HTTP GET (asgi.websocket.send)",
-            "HTTP GET (asgi.websocket.send)",
+            "/ asgi.websocket.receive",
+            "/ asgi.websocket.send",
+            "/ asgi.websocket.receive",
+            "/ asgi.websocket.send",
+            "/ asgi.websocket.receive",
+            "/ asgi",
         ]
         actual = [span.name for span in span_list]
-        self.assertListEqual(actual[:3], expected)
+        self.assertListEqual(actual, expected)
 
     def test_lifespan(self):
-        async def app_asgi(scope, receive, send):
-            assert scope["type"] == "lifespan"
-            await receive()
-            await send({})
-
         self.scope["type"] = "lifespan"
-        app = otel_asgi.OpenTelemetryMiddleware(app_asgi)
+        app = otel_asgi.OpenTelemetryMiddleware(simple_asgi)
         self.seed_app(app)
         self.send_default_request()
         span_list = self.memory_exporter.get_finished_spans()