Update test to use LogCapturer helper. Move shared test workflow and helpers to test/helpers

Author: VegetarianOrc

temporalio/contrib/opentelemetry.py

@@ -26,8 +26,7 @@
 import opentelemetry.trace.propagation.tracecontext
 import opentelemetry.util.types
 from opentelemetry.context import Context
-from opentelemetry.trace import Span, SpanKind, Status, StatusCode, _Links
-from opentelemetry.util import types
+from opentelemetry.trace import Status, StatusCode
 from typing_extensions import Protocol, TypeAlias, TypedDict
 
 import temporalio.activity
@@ -449,8 +448,7 @@ async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any:
             )
             return await super().handle_query(input)
         finally:
-            detach_context = opentelemetry.context.get_current()
-            if detach_context is attach_context:
+            if attach_context == opentelemetry.context.get_current():
                 opentelemetry.context.detach(token)
 
     def handle_update_validator(
@@ -541,8 +539,7 @@ def _top_level_workflow_context(
                     kind=opentelemetry.trace.SpanKind.INTERNAL,
                 )
 
-            detach_context = opentelemetry.context.get_current()
-            if detach_context is attach_context:
+            if attach_context == opentelemetry.context.get_current():
                 opentelemetry.context.detach(token)
 
     #

tests/contrib/test_opentelemetry.py

@@ -9,6 +9,7 @@
 from datetime import timedelta
 from typing import Iterable, List, Optional
 
+import opentelemetry.context
 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
@@ -22,11 +23,12 @@
 from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
-from tests.worker.test_workflow import (
+from tests.helpers.cache_evitction import (
     CacheEvictionTearDownWorkflow,
     WaitForeverWorkflow,
     wait_forever_activity,
 )
+from tests.helpers import LogCapturer
 
 # Passing through because Python 3.9 has an import bug at
 # https://github.com/python/cpython/issues/91351
@@ -471,19 +473,6 @@ async def test_opentelemetry_safe_detach(client: Client):
     provider.add_span_processor(SimpleSpanProcessor(exporter))
     tracer = get_tracer(__name__, tracer_provider=provider)
 
-    class _OtelLogSpy(logging.Handler):
-        def __init__(self, level: int | str = 0) -> None:
-            self.seenOtelFailedMessage = False
-            super().__init__(level)
-
-        def emit(self, record: logging.LogRecord) -> None:
-            if not self.seenOtelFailedMessage:
-                self.seenOtelFailedMessage = (
-                    record.levelno == logging.ERROR
-                    and record.name == "opentelemetry.context"
-                    and record.message == "Failed to detach context"
-                )
-
     async with Worker(
         client,
         workflows=[CacheEvictionTearDownWorkflow, WaitForeverWorkflow],
@@ -497,28 +486,34 @@ def emit(self, record: logging.LogRecord) -> None:
         old_hook = sys.unraisablehook
         hook_calls: List[sys.UnraisableHookArgs] = []
         sys.unraisablehook = hook_calls.append
-        log_spy = _OtelLogSpy()
-        logging.getLogger().addHandler(log_spy)
-        try:
-            handle = await client.start_workflow(
-                CacheEvictionTearDownWorkflow.run,
-                id=f"wf-{uuid.uuid4()}",
-                task_queue=worker.task_queue,
-            )
 
-            # CacheEvictionTearDownWorkflow requires 3 signals to be sent
-            await handle.signal(CacheEvictionTearDownWorkflow.signal)
-            await handle.signal(CacheEvictionTearDownWorkflow.signal)
-            await handle.signal(CacheEvictionTearDownWorkflow.signal)
+        with LogCapturer().logs_captured(opentelemetry.context.logger) as capturer:
+            try:
+                handle = await client.start_workflow(
+                    CacheEvictionTearDownWorkflow.run,
+                    id=f"wf-{uuid.uuid4()}",
+                    task_queue=worker.task_queue,
+                )
+
+                # CacheEvictionTearDownWorkflow requires 3 signals to be sent
+                await handle.signal(CacheEvictionTearDownWorkflow.signal)
+                await handle.signal(CacheEvictionTearDownWorkflow.signal)
+                await handle.signal(CacheEvictionTearDownWorkflow.signal)
 
-            await handle.result()
-        finally:
-            sys.unraisablehook = old_hook
-            logging.getLogger().removeHandler(log_spy)
+                await handle.result()
+            finally:
+                sys.unraisablehook = old_hook
 
-        # Confirm at least 1 exception
-        assert hook_calls
+            # Confirm at least 1 exception
+            assert hook_calls
+
+            def otel_context_error(record: logging.LogRecord) -> bool:
+                return (
+                    record.levelno == logging.ERROR
+                    and record.name == "opentelemetry.context"
+                    and record.message == "Failed to detach context"
+                )
 
-        assert (
-            not log_spy.seenOtelFailedMessage
-        ), "Detach from context message should not be logged"
+            assert (
+                capturer.find(otel_context_error) is None
+            ), "Detach from context message should not be logged"

tests/helpers/__init__.py

@@ -1,11 +1,25 @@
 import asyncio
+import logging
+import logging.handlers
+import queue
 import socket
 import time
 import uuid
-from contextlib import closing
+from contextlib import closing, contextmanager
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Any, Awaitable, Callable, Optional, Sequence, Type, TypeVar, Union
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from temporalio.api.common.v1 import WorkflowExecution
 from temporalio.api.enums.v1 import EventType as EventType
@@ -401,3 +415,34 @@ def _format_row(items: list[str], truncate: bool = False) -> str:
                 padding = len(f" *: {elapsed_ms:>4} ")
             summary_row[col_idx] = f"{' ' * padding}[{summary}]"[: col_width - 3]
             print(_format_row(summary_row))
+
+
+class LogCapturer:
+    def __init__(self) -> None:
+        self.log_queue: queue.Queue[logging.LogRecord] = queue.Queue()
+
+    @contextmanager
+    def logs_captured(self, *loggers: logging.Logger):
+        handler = logging.handlers.QueueHandler(self.log_queue)
+
+        prev_levels = [l.level for l in loggers]
+        for l in loggers:
+            l.setLevel(logging.INFO)
+            l.addHandler(handler)
+        try:
+            yield self
+        finally:
+            for i, l in enumerate(loggers):
+                l.removeHandler(handler)
+                l.setLevel(prev_levels[i])
+
+    def find_log(self, starts_with: str) -> Optional[logging.LogRecord]:
+        return self.find(lambda l: l.message.startswith(starts_with))
+
+    def find(
+        self, pred: Callable[[logging.LogRecord], bool]
+    ) -> Optional[logging.LogRecord]:
+        for record in cast(List[logging.LogRecord], self.log_queue.queue):
+            if pred(record):
+                return record
+        return None

tests/helpers/cache_evitction.py

@@ -0,0 +1,67 @@
+import asyncio
+from datetime import timedelta
+from temporalio import activity, workflow
+
+
+@activity.defn
+async def wait_forever_activity() -> None:
+    await asyncio.Future()
+
+
+@workflow.defn
+class WaitForeverWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        await asyncio.Future()
+
+
+@workflow.defn
+class CacheEvictionTearDownWorkflow:
+    def __init__(self) -> None:
+        self._signal_count = 0
+
+    @workflow.run
+    async def run(self) -> None:
+        # Start several things in background. This is just to show that eviction
+        # can work even with these things running.
+        tasks = [
+            asyncio.create_task(
+                workflow.execute_activity(
+                    wait_forever_activity, start_to_close_timeout=timedelta(hours=1)
+                )
+            ),
+            asyncio.create_task(
+                workflow.execute_child_workflow(WaitForeverWorkflow.run)
+            ),
+            asyncio.create_task(asyncio.sleep(1000)),
+            asyncio.shield(
+                workflow.execute_activity(
+                    wait_forever_activity, start_to_close_timeout=timedelta(hours=1)
+                )
+            ),
+            asyncio.create_task(workflow.wait_condition(lambda: False)),
+        ]
+        gather_fut = asyncio.gather(*tasks, return_exceptions=True)
+        # Let's also start something in the background that we never wait on
+        asyncio.create_task(asyncio.sleep(1000))
+        try:
+            # Wait for signal count to reach 2
+            await asyncio.sleep(0.01)
+            await workflow.wait_condition(lambda: self._signal_count > 1)
+        finally:
+            # This finally, on eviction, is actually called but the command
+            # should be ignored
+            await asyncio.sleep(0.01)
+            await workflow.wait_condition(lambda: self._signal_count > 2)
+            # Cancel gather tasks and wait on them, but ignore the errors
+            for task in tasks:
+                task.cancel()
+            await gather_fut
+
+    @workflow.signal
+    async def signal(self) -> None:
+        self._signal_count += 1
+
+    @workflow.query
+    def signal_count(self) -> int:
+        return self._signal_count

tests/worker/test_workflow.py

@@ -143,11 +143,17 @@
     pause_and_assert,
     unpause_and_assert,
     workflow_update_exists,
+    LogCapturer,
 )
 from tests.helpers.external_stack_trace import (
     ExternalStackTraceWorkflow,
     external_wait_cancel,
 )
+from tests.helpers.cache_evitction import (
+    CacheEvictionTearDownWorkflow,
+    WaitForeverWorkflow,
+    wait_forever_activity,
+)
 
 # Passing through because Python 3.9 has an import bug at
 # https://github.com/python/cpython/issues/91351
@@ -1996,37 +2002,6 @@ def last_signal(self) -> str:
         return self._last_signal
 
 
-class LogCapturer:
-    def __init__(self) -> None:
-        self.log_queue: queue.Queue[logging.LogRecord] = queue.Queue()
-
-    @contextmanager
-    def logs_captured(self, *loggers: logging.Logger):
-        handler = logging.handlers.QueueHandler(self.log_queue)
-
-        prev_levels = [l.level for l in loggers]
-        for l in loggers:
-            l.setLevel(logging.INFO)
-            l.addHandler(handler)
-        try:
-            yield self
-        finally:
-            for i, l in enumerate(loggers):
-                l.removeHandler(handler)
-                l.setLevel(prev_levels[i])
-
-    def find_log(self, starts_with: str) -> Optional[logging.LogRecord]:
-        return self.find(lambda l: l.message.startswith(starts_with))
-
-    def find(
-        self, pred: Callable[[logging.LogRecord], bool]
-    ) -> Optional[logging.LogRecord]:
-        for record in cast(List[logging.LogRecord], self.log_queue.queue):
-            if pred(record):
-                return record
-        return None
-
-
 async def test_workflow_logging(client: Client, env: WorkflowEnvironment):
     workflow.logger.full_workflow_info_on_extra = True
     with LogCapturer().logs_captured(
@@ -3739,70 +3714,6 @@ async def test_manual_result_type(client: Client):
         assert res4 == ManualResultType(some_string="from-query")
 
 
-@activity.defn
-async def wait_forever_activity() -> None:
-    await asyncio.Future()
-
-
-@workflow.defn
-class WaitForeverWorkflow:
-    @workflow.run
-    async def run(self) -> None:
-        await asyncio.Future()
-
-
-@workflow.defn
-class CacheEvictionTearDownWorkflow:
-    def __init__(self) -> None:
-        self._signal_count = 0
-
-    @workflow.run
-    async def run(self) -> None:
-        # Start several things in background. This is just to show that eviction
-        # can work even with these things running.
-        tasks = [
-            asyncio.create_task(
-                workflow.execute_activity(
-                    wait_forever_activity, start_to_close_timeout=timedelta(hours=1)
-                )
-            ),
-            asyncio.create_task(
-                workflow.execute_child_workflow(WaitForeverWorkflow.run)
-            ),
-            asyncio.create_task(asyncio.sleep(1000)),
-            asyncio.shield(
-                workflow.execute_activity(
-                    wait_forever_activity, start_to_close_timeout=timedelta(hours=1)
-                )
-            ),
-            asyncio.create_task(workflow.wait_condition(lambda: False)),
-        ]
-        gather_fut = asyncio.gather(*tasks, return_exceptions=True)
-        # Let's also start something in the background that we never wait on
-        asyncio.create_task(asyncio.sleep(1000))
-        try:
-            # Wait for signal count to reach 2
-            await asyncio.sleep(0.01)
-            await workflow.wait_condition(lambda: self._signal_count > 1)
-        finally:
-            # This finally, on eviction, is actually called but the command
-            # should be ignored
-            await asyncio.sleep(0.01)
-            await workflow.wait_condition(lambda: self._signal_count > 2)
-            # Cancel gather tasks and wait on them, but ignore the errors
-            for task in tasks:
-                task.cancel()
-            await gather_fut
-
-    @workflow.signal
-    async def signal(self) -> None:
-        self._signal_count += 1
-
-    @workflow.query
-    def signal_count(self) -> int:
-        return self._signal_count
-
-
 async def test_cache_eviction_tear_down(client: Client):
     # This test simulates forcing eviction. This used to raise GeneratorExit on
     # GC which triggered the finally which could run on any thread Python