wip: test reflection isn't broken

Author: gregbrowndev

tests/contrib/opentelemetry/helpers/reflection_interceptor.py

@@ -0,0 +1,77 @@
+import dataclasses
+import logging
+import typing
+
+import temporalio.worker
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class InterceptedActivity:
+    class_name: str
+    name: typing.Optional[str]
+    qualname: typing.Optional[str]
+    module: typing.Optional[str]
+    annotations: typing.Dict[str, typing.Any]
+    docstring: typing.Optional[str]
+
+
+class ReflectionInterceptor(temporalio.worker.Interceptor):
+    """Interceptor to check we haven't broken reflection when wrapping the activity."""
+
+    def __init__(self) -> None:
+        self._intercepted_activities: list[InterceptedActivity] = []
+
+    def get_intercepted_activities(self) -> typing.List[InterceptedActivity]:
+        """Get the list of intercepted activities."""
+        return self._intercepted_activities
+
+    def intercept_activity(
+        self, next: temporalio.worker.ActivityInboundInterceptor
+    ) -> temporalio.worker.ActivityInboundInterceptor:
+        """Method called for intercepting an activity.
+
+        Args:
+            next: The underlying inbound interceptor this interceptor should
+                delegate to.
+
+        Returns:
+            The new interceptor that will be used to for the activity.
+        """
+        return _ReflectionActivityInboundInterceptor(next, self)
+
+
+class _ReflectionActivityInboundInterceptor(
+    temporalio.worker.ActivityInboundInterceptor
+):
+    def __init__(
+        self,
+        next: temporalio.worker.ActivityInboundInterceptor,
+        root: ReflectionInterceptor,
+    ) -> None:
+        super().__init__(next)
+        self.root = root
+
+    async def execute_activity(
+        self, input: temporalio.worker.ExecuteActivityInput
+    ) -> typing.Any:
+        """Called to invoke the activity."""
+
+        try:
+            self.root._intercepted_activities.append(
+                InterceptedActivity(
+                    class_name=input.fn.__class__.__name__,
+                    name=getattr(input.fn, "__name__", None),
+                    qualname=getattr(input.fn, "__qualname__", None),
+                    module=getattr(input.fn, "__module__", None),
+                    docstring=getattr(input.fn, "__doc__", None),
+                    annotations=getattr(input.fn, "__annotations__", {}),
+                )
+            )
+        except AttributeError:
+            logger.exception(
+                "Activity function does not have expected attributes, skipping reflection."
+            )
+
+        return await self.next.execute_activity(input)

tests/contrib/opentelemetry/helpers/tracing.py

@@ -1,34 +1,18 @@
 from __future__ import annotations
 
-import asyncio
-import concurrent.futures
-import logging
 import multiprocessing
 import multiprocessing.managers
 import threading
 import typing
-import uuid
 from dataclasses import dataclass
-from datetime import timedelta
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
 
 import opentelemetry.trace
-from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
+from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import (
-    SimpleSpanProcessor,
     SpanExporter,
     SpanExportResult,
 )
-from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-from opentelemetry.trace import get_current_span, get_tracer
-
-from temporalio import activity, workflow
-from temporalio.client import Client
-from temporalio.common import RetryPolicy
-from temporalio.contrib.opentelemetry import TracingInterceptor
-from temporalio.contrib.opentelemetry import workflow as otel_workflow
-from temporalio.testing import WorkflowEnvironment
-from temporalio.worker import SharedStateManager, UnsandboxedWorkflowRunner, Worker
 
 
 @dataclass(frozen=True)
@@ -85,14 +69,9 @@ def from_readable_span(cls, span: ReadableSpan) -> "SerialisableSpan":
         )
 
 
-SerialisableSpanListProxy: typing.TypeAlias = multiprocessing.managers.ListProxy[
-    SerialisableSpan
-]
-
-
 def make_span_proxy_list(
     manager: multiprocessing.managers.SyncManager,
-) -> SerialisableSpanListProxy:
+) -> multiprocessing.managers.ListProxy[SerialisableSpan]:
     """Create a list proxy to share `SerialisableSpan` across processes."""
     return manager.list()
 
@@ -110,7 +89,9 @@ class _ListProxySpanExporter(SpanExporter):
     into a single trace.
     """
 
-    def __init__(self, finished_spans: SerialisableSpanListProxy) -> None:
+    def __init__(
+        self, finished_spans: multiprocessing.managers.ListProxy[SerialisableSpan]
+    ) -> None:
         self._finished_spans = finished_spans
         self._stopped = False
         self._lock = threading.Lock()

tests/contrib/test_opentelemetry.py

@@ -26,12 +26,15 @@
 from temporalio.contrib.opentelemetry import workflow as otel_workflow
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import SharedStateManager, UnsandboxedWorkflowRunner, Worker
+from tests.contrib.opentelemetry.helpers.reflection_interceptor import (
+    InterceptedActivity,
+    ReflectionInterceptor,
+)
 from tests.contrib.opentelemetry.helpers.tracing import (
     SerialisableSpan,
     _ListProxySpanExporter,
     dump_spans,
     make_span_proxy_list,
-    SerialisableSpanListProxy,
 )
 
 # Passing through because Python 3.9 has an import bug at
@@ -403,9 +406,6 @@ async def test_activity_trace_propagation(
     client: Client,
     env: WorkflowEnvironment,
 ):
-    # TODO: add spy interceptor to check `input.fn` wraps original metadata
-    # TODO: Add Resource to show how resource would be propagated
-
     # Create a tracer that has an in-memory exporter
     exporter = InMemorySpanExporter()
     provider = TracerProvider()
@@ -417,13 +417,16 @@ async def test_activity_trace_propagation(
     manager = multiprocessing.Manager()
     finished_spans_proxy = make_span_proxy_list(manager)
 
+    # Create an interceptor to test we haven't broken reflection
+    reflection_interceptor = ReflectionInterceptor()
+
     # Create a worker with a process pool activity executor
     async with Worker(
         client,
         task_queue=f"task_queue_{uuid.uuid4()}",
         workflows=[ActivityTracePropagationWorkflow],
         activities=[sync_activity],
-        interceptors=[TracingInterceptor(tracer)],
+        interceptors=[TracingInterceptor(tracer), reflection_interceptor],
         activity_executor=concurrent.futures.ProcessPoolExecutor(
             max_workers=1,
             initializer=activity_trace_propagation_initializer,
@@ -437,16 +440,29 @@ async def test_activity_trace_propagation(
             task_queue=worker.task_queue,
         )
 
+    # The dumped spans should include child spans created in the child process
     spans = exporter.get_finished_spans() + tuple(finished_spans_proxy)
     logging.debug("Spans:\n%s", "\n".join(dump_spans(spans, with_attributes=False)))
     assert dump_spans(spans, with_attributes=False) == [
         "RunActivity:sync_activity",
         "  child_span",
     ]
 
+    # and the activity should still have the original attributes in downstream interceptors
+    assert reflection_interceptor.get_intercepted_activities() == [
+        InterceptedActivity(
+            class_name="ActivityFnWithTraceContext",
+            name="sync_activity",
+            qualname="sync_activity",
+            module="tests.contrib.test_opentelemetry",
+            docstring="An activity that uses tracing features.\n\nWhen executed in a process pool, we expect the trace context to be available\nfrom the parent process.\n",
+            annotations={"param": "typing.Any", "return": "str"},
+        )
+    ]
+
 
 def activity_trace_propagation_initializer(
-    _finished_spans_proxy: SerialisableSpanListProxy,
+    _finished_spans_proxy: multiprocessing.managers.ListProxy[SerialisableSpan],
 ) -> None:
     """Initializer for the process pool worker to export spans to a shared list."""
     _exporter = _ListProxySpanExporter(_finished_spans_proxy)