Skip to content

Commit 89dea67

Browse files
goingforstudying-ctrlgoingforstudying-ctrl
authored andcommitted
fix: use create_autospec(Transport) and replace caplog with _logger mock
Address review feedback: - Use mock.create_autospec(Transport) per henry3260's suggestion - Replace caplog with mock.patch on _logger to avoid logging config modification
1 parent d0a2f1a commit 89dea67

1 file changed

Lines changed: 7 additions & 5 deletions

File tree

providers/google/tests/unit/google/cloud/log/test_stackdriver_task_handler.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
import pytest
2727
from google.cloud.logging import Resource
28+
from google.cloud.logging.handlers.transports import Transport
2829
from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse, LogEntry
2930

3031
from airflow.providers.google.cloud.log.stackdriver_task_handler import (
@@ -326,9 +327,10 @@ def test_processors_skips_non_task_logger(self, mock_client, mock_get_creds_and_
326327
mock_transport_type.return_value.send.assert_not_called()
327328

328329
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="airflow.sdk.log only exists in Airflow 3+")
330+
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler._logger")
329331
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id")
330332
@mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client")
331-
def test_processors_survives_transport_send_failure(self, mock_client, mock_get_creds_and_project_id, caplog):
333+
def test_processors_survives_transport_send_failure(self, mock_client, mock_get_creds_and_project_id, mock_logger):
332334
"""``proc()`` must not let a transport.send() failure crash the process.
333335
334336
In AF3's supervisor model REMOTE_TASK_LOG applies to ALL supervised
@@ -339,7 +341,7 @@ def test_processors_survives_transport_send_failure(self, mock_client, mock_get_
339341
"""
340342
mock_get_creds_and_project_id.return_value = ("creds", "project_id")
341343

342-
mock_transport_type = mock.MagicMock()
344+
mock_transport_type = mock.create_autospec(Transport)
343345
mock_transport_type.return_value.send.side_effect = RuntimeError("IAM permission denied")
344346
with mock.patch("airflow.sdk.log.relative_path_from_logger", return_value="dag/task/1.log"):
345347
io = StackdriverRemoteLogIO(
@@ -354,13 +356,13 @@ def test_processors_survives_transport_send_failure(self, mock_client, mock_get_
354356
"logger_name": "airflow.task",
355357
"timestamp": "2026-01-15T10:30:00+00:00",
356358
}
357-
with caplog.at_level(logging.WARNING):
358-
result = proc(mock.MagicMock(), "info", event)
359+
result = proc(mock.MagicMock(), "info", event)
359360

360361
# Must still return the event — crash would mean no return at all
361362
assert result is event
362363
# Must log the failure so the operator can see it
363-
assert "Failed to send log to Cloud Logging" in caplog.text
364+
mock_logger.warning.assert_called_once()
365+
assert "Failed to send log to Cloud Logging" in mock_logger.warning.call_args[0][0]
364366

365367

366368
@pytest.mark.usefixtures("clean_stackdriver_handlers")

0 commit comments

Comments
 (0)