tests(worker): consolidate workflow logging test; merge trace identifier into task-failure; restore logger state in finally

Author: harsh543

tests/worker/test_workflow.py

@@ -1993,12 +1993,16 @@ def last_signal(self) -> str:
 
 @pytest.mark.parametrize("temporal_extra_mode", ["dict", "flatten"])
 async def test_workflow_logging(client: Client, temporal_extra_mode: str):
-    """Test that workflow logger produces correct log records for each extra mode."""
+    """Test workflow logging: extra mode formatting, replay suppression, and full_workflow_info."""
     original_mode = workflow.logger.temporal_extra_mode
+    original_full_info = workflow.logger.full_workflow_info_on_extra
     workflow.logger.temporal_extra_mode = temporal_extra_mode
+    workflow.logger.full_workflow_info_on_extra = True
 
     try:
         with LogCapturer().logs_captured(workflow.logger.base_logger) as capturer:
+            # --- First execution: logs should appear ---
+            # Disable workflow cache so worker restart triggers replay
             async with new_worker(
                 client, LoggingWorkflow, max_cached_workflows=0
             ) as worker:
@@ -2008,19 +2012,26 @@ async def test_workflow_logging(client: Client, temporal_extra_mode: str):
                     task_queue=worker.task_queue,
                 )
                 await handle.signal(LoggingWorkflow.my_signal, "signal 1")
+                await handle.signal(LoggingWorkflow.my_signal, "signal 2")
                 await handle.execute_update(
                     LoggingWorkflow.my_update, "update 1", id="update-1"
                 )
-                await handle.signal(LoggingWorkflow.my_signal, "finish")
-                await handle.result()
+                assert "signal 2" == await handle.query(LoggingWorkflow.last_signal)
 
+            # Verify logs from first execution
             record = capturer.find_log("Signal: signal 1")
             assert record is not None
             assert record.funcName == "my_signal"
+            assert capturer.find_log("Signal: signal 2")
+            assert capturer.find_log("Query called")
 
             update_record = capturer.find_log("Update: update 1")
             assert update_record is not None
 
+            # Verify full_workflow_info_on_extra
+            assert isinstance(record.__dict__["workflow_info"], workflow.Info)
+
+            # Verify extra mode formatting
             if temporal_extra_mode == "dict":
                 # Dict mode appends context to message and uses nested dict
                 assert "({'attempt':" in record.message
@@ -2046,49 +2057,6 @@ async def test_workflow_logging(client: Client, temporal_extra_mode: str):
                         assert isinstance(
                             value, (str, int, float, bool, type(None))
                         ), f"Key {key} has non-primitive value: {type(value)}"
-    finally:
-        workflow.logger.temporal_extra_mode = original_mode
-
-
-async def test_workflow_logging_replay(client: Client):
-    """Test replay log suppression and full_workflow_info_on_extra.
-
-    This test validates behavior the parameterized dict/flatten test does NOT cover:
-    1. Replay suppression: logs from replayed history should not be emitted again.
-    2. full_workflow_info_on_extra: when enabled, logs include workflow.Info object.
-    """
-    original_full_info = workflow.logger.full_workflow_info_on_extra
-    workflow.logger.full_workflow_info_on_extra = True
-    try:
-        with LogCapturer().logs_captured(
-            workflow.logger.base_logger, activity.logger.base_logger
-        ) as capturer:
-            # --- First execution: logs should appear ---
-            # Disable workflow cache so worker restart triggers replay
-            async with new_worker(
-                client, LoggingWorkflow, max_cached_workflows=0
-            ) as worker:
-                handle = await client.start_workflow(
-                    LoggingWorkflow.run,
-                    id=f"workflow-{uuid.uuid4()}",
-                    task_queue=worker.task_queue,
-                )
-                await handle.signal(LoggingWorkflow.my_signal, "signal 1")
-                await handle.signal(LoggingWorkflow.my_signal, "signal 2")
-                await handle.execute_update(
-                    LoggingWorkflow.my_update, "update 1", id="update-1"
-                )
-                assert "signal 2" == await handle.query(LoggingWorkflow.last_signal)
-
-            # Verify logs from first execution
-            assert capturer.find_log("Signal: signal 1")
-            assert capturer.find_log("Signal: signal 2")
-            assert capturer.find_log("Update: update 1")
-            assert capturer.find_log("Query called")
-
-            # Verify full_workflow_info_on_extra
-            record = capturer.find_log("Signal: signal 1")
-            assert record and isinstance(record.__dict__["workflow_info"], workflow.Info)
 
             # --- Clear logs and continue execution (replay path) ---
             # When the new worker starts, it replays the workflow history (signals 1 & 2).
@@ -2111,6 +2079,7 @@ async def test_workflow_logging_replay(client: Client):
             assert capturer.find_log("Signal: signal 3")
             assert capturer.find_log("Signal: finish")
     finally:
+        workflow.logger.temporal_extra_mode = original_mode
         workflow.logger.full_workflow_info_on_extra = original_full_info
 
 
@@ -2175,6 +2144,18 @@ async def test_workflow_logging_task_fail(client: Client):
             == "task_fail_once_activity"
         )
 
+        def workflow_failure_with_identifier(l: logging.LogRecord):
+            if (
+                hasattr(l, "__temporal_error_identifier")
+                and getattr(l, "__temporal_error_identifier")
+                == "WorkflowTaskFailure"
+            ):
+                assert l.msg.startswith("Failed activation on workflow")
+                return True
+            return False
+
+        assert capturer.find(workflow_failure_with_identifier) is not None
+
 
 @workflow.defn
 class StackTraceWorkflow:
@@ -8034,33 +8015,6 @@ async def test_quick_activity_swallows_cancellation(client: Client):
             assert cause.message == "Workflow cancelled"
 
 
-async def test_workflow_logging_trace_identifier(client: Client):
-    with LogCapturer().logs_captured(
-        temporalio.worker._workflow_instance.logger
-    ) as capturer:
-        async with new_worker(
-            client,
-            TaskFailOnceWorkflow,
-            activities=[task_fail_once_activity],
-        ) as worker:
-            await client.execute_workflow(
-                TaskFailOnceWorkflow.run,
-                id="workflow_failure_trace_identifier",
-                task_queue=worker.task_queue,
-            )
-
-        def workflow_failure(l: logging.LogRecord):
-            if (
-                hasattr(l, "__temporal_error_identifier")
-                and getattr(l, "__temporal_error_identifier") == "WorkflowTaskFailure"
-            ):
-                assert l.msg.startswith("Failed activation on workflow")
-                return True
-            return False
-
-        assert capturer.find(workflow_failure) is not None
-
-
 @activity.defn
 def use_in_workflow() -> bool:
     return workflow.in_workflow()