Context manager for logger handlers in tests (#1293)

Author: jmaeagle99

tests/helpers/__init__.py

@@ -5,7 +5,7 @@
 import socket
 import time
 import uuid
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Awaitable, Callable, Iterator, Sequence
 from contextlib import closing, contextmanager
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
@@ -440,3 +440,16 @@ def find(
             if pred(record):
                 return record
         return None
+
+
+class LogHandler:
+    @staticmethod
+    @contextmanager
+    def apply(logger: logging.Logger, handler: logging.Handler) -> Iterator[None]:
+        level = logger.level
+        logger.addHandler(handler)
+        try:
+            yield
+        finally:
+            logger.removeHandler(handler)
+            logger.level = level

tests/test_runtime.py

@@ -22,6 +22,7 @@
 )
 from temporalio.worker import Worker
 from tests.helpers import (
+    LogHandler,
     assert_eq_eventually,
     assert_eventually,
     find_free_port,
@@ -80,8 +81,8 @@ async def test_runtime_log_forwarding():
     # Create logger with record capture
     log_queue: queue.Queue[logging.LogRecord] = queue.Queue()
     log_queue_list = cast(list[logging.LogRecord], log_queue.queue)
+    handler = logging.handlers.QueueHandler(log_queue)
     logger = logging.getLogger(f"log-{uuid.uuid4()}")
-    logger.addHandler(logging.handlers.QueueHandler(log_queue))
 
     async def log_queue_len() -> int:
         return len(log_queue_list)
@@ -96,49 +97,50 @@ async def log_queue_len() -> int:
         )
     )
 
-    # Set capture only info logs
-    logger.setLevel(logging.INFO)
-    # Write some logs
-    runtime._core_runtime.write_test_info_log("info1", "extra1")
-    runtime._core_runtime.write_test_debug_log("debug2", "extra2")
-    runtime._core_runtime.write_test_info_log("info3", "extra3")
-
-    # Check the expected records
-    await assert_eq_eventually(2, log_queue_len)
-    assert log_queue_list[0].levelno == logging.INFO
-    assert log_queue_list[0].message.startswith(
-        "[sdk_core::temporal_sdk_bridge::runtime] info1"
-    )
-    assert (
-        log_queue_list[0].name
-        == f"{logger.name}-sdk_core::temporal_sdk_bridge::runtime"
-    )
-    assert log_queue_list[0].created == log_queue_list[0].temporal_log.time  # type: ignore
-    assert log_queue_list[0].temporal_log.fields == {"extra_data": "extra1"}  # type: ignore
-    assert log_queue_list[1].levelno == logging.INFO
-    assert log_queue_list[1].message.startswith(
-        "[sdk_core::temporal_sdk_bridge::runtime] info3"
-    )
+    with LogHandler.apply(logger, handler):
+        # Set capture only info logs
+        logger.setLevel(logging.INFO)
+        # Write some logs
+        runtime._core_runtime.write_test_info_log("info1", "extra1")
+        runtime._core_runtime.write_test_debug_log("debug2", "extra2")
+        runtime._core_runtime.write_test_info_log("info3", "extra3")
+
+        # Check the expected records
+        await assert_eq_eventually(2, log_queue_len)
+        assert log_queue_list[0].levelno == logging.INFO
+        assert log_queue_list[0].message.startswith(
+            "[sdk_core::temporal_sdk_bridge::runtime] info1"
+        )
+        assert (
+            log_queue_list[0].name
+            == f"{logger.name}-sdk_core::temporal_sdk_bridge::runtime"
+        )
+        assert log_queue_list[0].created == log_queue_list[0].temporal_log.time  # type: ignore
+        assert log_queue_list[0].temporal_log.fields == {"extra_data": "extra1"}  # type: ignore
+        assert log_queue_list[1].levelno == logging.INFO
+        assert log_queue_list[1].message.startswith(
+            "[sdk_core::temporal_sdk_bridge::runtime] info3"
+        )
 
-    # Clear logs and enable debug and try again
-    log_queue_list.clear()
-    logger.setLevel(logging.DEBUG)
-    runtime._core_runtime.write_test_info_log("info4", "extra4")
-    runtime._core_runtime.write_test_debug_log("debug5", "extra5")
-    runtime._core_runtime.write_test_info_log("info6", "extra6")
-    await assert_eq_eventually(3, log_queue_len)
-    assert log_queue_list[0].levelno == logging.INFO
-    assert log_queue_list[0].message.startswith(
-        "[sdk_core::temporal_sdk_bridge::runtime] info4"
-    )
-    assert log_queue_list[1].levelno == logging.DEBUG
-    assert log_queue_list[1].message.startswith(
-        "[sdk_core::temporal_sdk_bridge::runtime] debug5"
-    )
-    assert log_queue_list[2].levelno == logging.INFO
-    assert log_queue_list[2].message.startswith(
-        "[sdk_core::temporal_sdk_bridge::runtime] info6"
-    )
+        # Clear logs and enable debug and try again
+        log_queue_list.clear()
+        logger.setLevel(logging.DEBUG)
+        runtime._core_runtime.write_test_info_log("info4", "extra4")
+        runtime._core_runtime.write_test_debug_log("debug5", "extra5")
+        runtime._core_runtime.write_test_info_log("info6", "extra6")
+        await assert_eq_eventually(3, log_queue_len)
+        assert log_queue_list[0].levelno == logging.INFO
+        assert log_queue_list[0].message.startswith(
+            "[sdk_core::temporal_sdk_bridge::runtime] info4"
+        )
+        assert log_queue_list[1].levelno == logging.DEBUG
+        assert log_queue_list[1].message.startswith(
+            "[sdk_core::temporal_sdk_bridge::runtime] debug5"
+        )
+        assert log_queue_list[2].levelno == logging.INFO
+        assert log_queue_list[2].message.startswith(
+            "[sdk_core::temporal_sdk_bridge::runtime] info6"
+        )
 
 
 @workflow.defn
@@ -152,8 +154,8 @@ async def test_runtime_task_fail_log_forwarding(client: Client):
     # Client with lo capturing runtime
     log_queue: queue.Queue[logging.LogRecord] = queue.Queue()
     log_queue_list = cast(list[logging.LogRecord], log_queue.queue)
+    handler = logging.handlers.QueueHandler(log_queue)
     logger = logging.getLogger(f"log-{uuid.uuid4()}")
-    logger.addHandler(logging.handlers.QueueHandler(log_queue))
     logger.setLevel(logging.WARN)
     client = await Client.connect(
         client.service_client.config.target_host,
@@ -168,30 +170,32 @@ async def test_runtime_task_fail_log_forwarding(client: Client):
         ),
     )
 
-    # Start workflow
-    task_queue = f"task-queue-{uuid.uuid4()}"
-    async with Worker(client, task_queue=task_queue, workflows=[TaskFailWorkflow]):
-        handle = await client.start_workflow(
-            TaskFailWorkflow.run,
-            id=f"workflow-{uuid.uuid4()}",
-            task_queue=task_queue,
-        )
-
-        # Wait for log to appear
-        async def has_log() -> bool:
-            return any(
-                l for l in log_queue_list if "Failing workflow task" in l.message
+    with LogHandler.apply(logger, handler):
+        # Start workflow
+        task_queue = f"task-queue-{uuid.uuid4()}"
+        async with Worker(client, task_queue=task_queue, workflows=[TaskFailWorkflow]):
+            handle = await client.start_workflow(
+                TaskFailWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=task_queue,
             )
 
-        await assert_eq_eventually(True, has_log)
+            # Wait for log to appear
+            async def has_log() -> bool:
+                return any(
+                    l for l in log_queue_list if "Failing workflow task" in l.message
+                )
 
-    # Check record
-    record = next(l for l in log_queue_list if "Failing workflow task" in l.message)
-    assert record.levelno == logging.WARNING
-    assert (
-        record.name == f"{logger.name}-sdk_core::temporalio_sdk_core::worker::workflow"
-    )
-    assert record.temporal_log.fields["run_id"] == handle.result_run_id  # type: ignore
+            await assert_eq_eventually(True, has_log)
+
+        # Check record
+        record = next(l for l in log_queue_list if "Failing workflow task" in l.message)
+        assert record.levelno == logging.WARNING
+        assert (
+            record.name
+            == f"{logger.name}-sdk_core::temporalio_sdk_core::worker::workflow"
+        )
+        assert record.temporal_log.fields["run_id"] == handle.result_run_id  # type: ignore
 
 
 async def test_prometheus_histogram_bucket_overrides(client: Client):

tests/worker/test_activity.py

@@ -46,6 +46,7 @@
     Worker,
     WorkerConfig,
 )
+from tests.helpers import LogHandler
 from tests.helpers.worker import (
     ExternalWorker,
     KSAction,
@@ -1019,20 +1020,15 @@ async def say_hello(name: str) -> str:
 
     # Create a queue, add handler to logger, call normal activity, then check
     handler = logging.handlers.QueueHandler(queue.Queue())
-    activity.logger.base_logger.addHandler(handler)
-    prev_level = activity.logger.base_logger.level
-    activity.logger.base_logger.setLevel(logging.INFO)
-    try:
+    with LogHandler.apply(activity.logger.base_logger, handler):
+        activity.logger.base_logger.setLevel(logging.INFO)
         result = await _execute_workflow_with_activity(
             client,
             worker,
             say_hello,
             "Temporal",
             shared_state_manager=shared_state_manager,
         )
-    finally:
-        activity.logger.base_logger.removeHandler(handler)
-        activity.logger.base_logger.setLevel(prev_level)
     assert result.result == "Hello, Temporal!"
     records: list[logging.LogRecord] = list(handler.queue.queue)  # type: ignore
     assert len(records) > 0
@@ -1671,9 +1667,8 @@ async def raise_error():
         raise RuntimeError("oh no!")
 
     handler = CustomLogHandler()
-    activity.logger.base_logger.addHandler(handler)
 
-    try:
+    with LogHandler.apply(activity.logger.base_logger, handler):
         with pytest.raises(WorkflowFailureError) as err:
             await _execute_workflow_with_activity(
                 client,
@@ -1686,9 +1681,6 @@ async def raise_error():
         )
         assert handler._trace_identifiers == 1
 
-    finally:
-        activity.logger.base_logger.removeHandler(CustomLogHandler())
-
 
 async def test_activity_heartbeat_context(
     client: Client,

tests/worker/test_workflow.py

@@ -122,6 +122,7 @@
 from tests import DEV_SERVER_DOWNLOAD_VERSION
 from tests.helpers import (
     LogCapturer,
+    LogHandler,
     admitted_update_task,
     assert_eq_eventually,
     assert_eventually,
@@ -8445,41 +8446,42 @@ async def run(self):
 
 class CustomLogHandler(logging.Handler):
     def emit(self, record: logging.LogRecord) -> None:
-        import httpx  # type: ignore[reportUnusedImport]
+        import httpx  # type: ignore[reportUnusedImport] # noqa
 
 
 async def test_disable_logger_sandbox(
     client: Client,
 ):
     logger = workflow.logger.logger
-    logger.addHandler(CustomLogHandler())
-    async with new_worker(
-        client,
-        DisableLoggerSandbox,
-        activities=[],
-    ) as worker:
-        with pytest.raises(WorkflowFailureError):
-            await client.execute_workflow(
-                DisableLoggerSandbox.run,
-                id=f"workflow-{uuid.uuid4()}",
-                task_queue=worker.task_queue,
-                run_timeout=timedelta(seconds=1),
-                retry_policy=RetryPolicy(maximum_attempts=1),
-            )
-        workflow.logger.unsafe_disable_sandbox()
-        await client.execute_workflow(
-            DisableLoggerSandbox.run,
-            id=f"workflow-{uuid.uuid4()}",
-            task_queue=worker.task_queue,
-            run_timeout=timedelta(seconds=1),
-            retry_policy=RetryPolicy(maximum_attempts=1),
-        )
-        workflow.logger.unsafe_disable_sandbox(False)
-        with pytest.raises(WorkflowFailureError):
+    handler = CustomLogHandler()
+    with LogHandler.apply(logger, handler):
+        async with new_worker(
+            client,
+            DisableLoggerSandbox,
+            activities=[],
+        ) as worker:
+            with pytest.raises(WorkflowFailureError):
+                await client.execute_workflow(
+                    DisableLoggerSandbox.run,
+                    id=f"workflow-{uuid.uuid4()}",
+                    task_queue=worker.task_queue,
+                    run_timeout=timedelta(seconds=1),
+                    retry_policy=RetryPolicy(maximum_attempts=1),
+                )
+            workflow.logger.unsafe_disable_sandbox()
             await client.execute_workflow(
                 DisableLoggerSandbox.run,
                 id=f"workflow-{uuid.uuid4()}",
                 task_queue=worker.task_queue,
                 run_timeout=timedelta(seconds=1),
                 retry_policy=RetryPolicy(maximum_attempts=1),
             )
+            workflow.logger.unsafe_disable_sandbox(False)
+            with pytest.raises(WorkflowFailureError):
+                await client.execute_workflow(
+                    DisableLoggerSandbox.run,
+                    id=f"workflow-{uuid.uuid4()}",
+                    task_queue=worker.task_queue,
+                    run_timeout=timedelta(seconds=1),
+                    retry_policy=RetryPolicy(maximum_attempts=1),
+                )