From 33b1002bfba0b99140771df9411c8de04ca7bfe2 Mon Sep 17 00:00:00 2001
From: Becky Smith <becky.smith@thedatalab.org>
Date: Fri, 1 Nov 2024 16:30:40 +0000
Subject: [PATCH] Add tracing tests

---
 tests/integration/test_middleware.py | 24 ++++++++++++++++++++++++
 tests/test_tracing.py                | 17 +++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tests/integration/test_middleware.py b/tests/integration/test_middleware.py
index fcb63974..f3bf15be 100644
--- a/tests/integration/test_middleware.py
+++ b/tests/integration/test_middleware.py
@@ -3,7 +3,10 @@
 import pytest
 from django.conf import settings
 
+from opentelemetry import trace
+
 from tests import factories
+from tests.conftest import get_trace
 
 
 @pytest.mark.django_db
@@ -54,3 +57,24 @@ def test_middleware_expired_error(airlock_client, responses):
 
     response = airlock_client.get("/workspaces/")
     assert response.status_code == 200
+
+
+@pytest.mark.django_db
+def test_middleware_user_trace(airlock_client, responses):
+    user = factories.create_user(workspaces=["workspace"])
+    airlock_client.login_with_user(user)
+    factories.create_workspace("workspace")
+
+    # In tests the current span in the middleware is a NonRecordingSpan,
+    # so call the endpoint inside another span so we can assert that the
+    # user is added during the middleware
+    tracer = trace.get_tracer("test")
+    with tracer.start_as_current_span("mock_django_span"):
+        response = airlock_client.get("/workspaces/view/workspace/")
+    
+    assert response.status_code == 200
+
+    mock_django_span_attributes, = [
+        span.attributes for span in get_trace() if span.name == "mock_django_span"
+    ]
+    assert mock_django_span_attributes == {"workspace": "workspace", "user": user.username}
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
index fb3432a6..7e036214 100644
--- a/tests/test_tracing.py
+++ b/tests/test_tracing.py
@@ -8,6 +8,7 @@
 from opentelemetry.sdk.trace.export import ConsoleSpanExporter
 
 from services.tracing import instrument, setup_default_tracing
+from tests.conftest import get_trace
 
 
 def test_setup_default_tracing_empty_env(monkeypatch):
@@ -83,6 +84,22 @@ def test_instrument_decorator_with_name_and_attributes():
     assert current_span.attributes == {"foo": "bar"}  # type: ignore
 
 
+def test_instrument_decorator_parent_attributes(settings):
+    @instrument(span_name="child", attributes={"foo": "bar"})
+    def child(): ...
+
+    tracer = trace.get_tracer("test")
+    with tracer.start_as_current_span("parent", attributes={"p_foo": "p_bar"}):
+        child()
+
+    spans = {span.name: span.attributes for span in get_trace()}
+
+    assert spans == {
+        "parent": {"p_foo": "p_bar", "foo": "bar"},
+        "child": {"foo": "bar"},
+    }
+
+
 @pytest.mark.parametrize(
     "func_attributes,func_args,func_kwargs,expected_attributes",
     [