Remove logger handlers in tests

Author: jmaeagle99

tests/test_runtime.py

@@ -80,8 +80,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 +96,53 @@ 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"
-    )
+    logger.addHandler(handler)
+    try:
+        # 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"
+        )
+    finally:
+        logger.removeHandler(handler)
 
 
 @workflow.defn
@@ -152,8 +156,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 +172,35 @@ 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
+    logger.addHandler(handler)
+    try:
+        # 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
+    finally:
+        logger.removeHandler(handler)
 
 
 async def test_prometheus_histogram_bucket_overrides(client: Client):

tests/worker/test_activity.py

@@ -1687,7 +1687,7 @@ async def raise_error():
         assert handler._trace_identifiers == 1
 
     finally:
-        activity.logger.base_logger.removeHandler(CustomLogHandler())
+        activity.logger.base_logger.removeHandler(handler)
 
 
 async def test_activity_heartbeat_context(

tests/worker/test_workflow.py

@@ -8445,41 +8445,45 @@ 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()
+    logger.addHandler(handler)
+    try:
+        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),
+                )
+    finally:
+        logger.removeHandler(handler)