Change to custom context manager

Author: jmaeagle99

tests/helpers/__init__.py

@@ -11,6 +11,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import (
     Any,
+    Iterator,
     TypeVar,
     cast,
 )
@@ -440,3 +441,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,
@@ -96,8 +97,7 @@ async def log_queue_len() -> int:
         )
     )
 
-    logger.addHandler(handler)
-    try:
+    with LogHandler.apply(logger, handler):
         # Set capture only info logs
         logger.setLevel(logging.INFO)
         # Write some logs
@@ -141,8 +141,6 @@ async def log_queue_len() -> int:
         assert log_queue_list[2].message.startswith(
             "[sdk_core::temporal_sdk_bridge::runtime] info6"
         )
-    finally:
-        logger.removeHandler(handler)
 
 
 @workflow.defn
@@ -172,8 +170,7 @@ async def test_runtime_task_fail_log_forwarding(client: Client):
         ),
     )
 
-    logger.addHandler(handler)
-    try:
+    with LogHandler.apply(logger, handler):
         # Start workflow
         task_queue = f"task-queue-{uuid.uuid4()}"
         async with Worker(client, task_queue=task_queue, workflows=[TaskFailWorkflow]):
@@ -199,8 +196,6 @@ async def has_log() -> bool:
             == 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

@@ -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(handler)
-
 
 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,
@@ -8453,8 +8454,7 @@ async def test_disable_logger_sandbox(
 ):
     logger = workflow.logger.logger
     handler = CustomLogHandler()
-    logger.addHandler(handler)
-    try:
+    with LogHandler.apply(logger, handler):
         async with new_worker(
             client,
             DisableLoggerSandbox,
@@ -8485,5 +8485,3 @@ async def test_disable_logger_sandbox(
                     run_timeout=timedelta(seconds=1),
                     retry_policy=RetryPolicy(maximum_attempts=1),
                 )
-    finally:
-        logger.removeHandler(handler)