(airflow) 01/01: Restore airflow_local_settings after the test_should_use_configured_log_name

[ta...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

taragolis pushed a commit to branch restore-local-settings-after-stackdriver-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit b28875960c65359982231677c2f8b07922debe49
Author: Andrey Anshin <An...@taragol.is>
AuthorDate: Wed Apr 3 21:22:15 2024 +0400

    Restore airflow_local_settings after the test_should_use_configured_log_name
---
 .../cloud/log/test_stackdriver_task_handler.py     | 40 +++++++++++++---------
 1 file changed, 23 insertions(+), 17 deletions(-)

diff --git a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py
index 5c01b9ff52..737cedefaa 100644
--- a/tests/providers/google/cloud/log/test_stackdriver_task_handler.py
+++ b/tests/providers/google/cloud/log/test_stackdriver_task_handler.py
@@ -27,6 +27,7 @@ from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesR
 from airflow.providers.google.cloud.log.stackdriver_task_handler import StackdriverTaskHandler
 from airflow.utils import timezone
 from airflow.utils.state import TaskInstanceState
+from tests.test_utils.config import conf_vars
 from tests.test_utils.db import clear_db_dags, clear_db_runs
 
 
@@ -71,29 +72,34 @@ def test_should_pass_message_to_client(mock_client, mock_get_creds_and_project_i
 @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.get_credentials_and_project_id")
 @mock.patch("airflow.providers.google.cloud.log.stackdriver_task_handler.gcp_logging.Client")
 def test_should_use_configured_log_name(mock_client, mock_get_creds_and_project_id):
-    mock_get_creds_and_project_id.return_value = ("creds", "project_id")
+    import importlib
+    import logging
 
-    with mock.patch.dict(
-        "os.environ",
-        AIRFLOW__LOGGING__REMOTE_LOGGING="true",
-        AIRFLOW__LOGGING__REMOTE_BASE_LOG_FOLDER="stackdriver://host/path",
-    ):
-        import importlib
-        import logging
+    from airflow import settings
+    from airflow.config_templates import airflow_local_settings
 
-        from airflow import settings
-        from airflow.config_templates import airflow_local_settings
+    mock_get_creds_and_project_id.return_value = ("creds", "project_id")
 
+    try:
+        with conf_vars(
+            {
+                ("logging", "remote_logging"): "True",
+                ("logging", "remote_base_log_folder"): "stackdriver://host/path",
+            }
+        ):
+            importlib.reload(airflow_local_settings)
+            settings.configure_logging()
+
+            logger = logging.getLogger("airflow.task")
+            handler = logger.handlers[0]
+            assert isinstance(handler, StackdriverTaskHandler)
+            with mock.patch.object(handler, "transport_type") as transport_type_mock:
+                logger.error("foo")
+                transport_type_mock.assert_called_once_with(mock_client.return_value, "path")
+    finally:
         importlib.reload(airflow_local_settings)
         settings.configure_logging()
 
-        logger = logging.getLogger("airflow.task")
-        handler = logger.handlers[0]
-        assert isinstance(handler, StackdriverTaskHandler)
-        with mock.patch.object(handler, "transport_type") as transport_type_mock:
-            logger.error("foo")
-            transport_type_mock.assert_called_once_with(mock_client.return_value, "path")
-
 
 @pytest.mark.db_test
 class TestStackdriverLoggingHandlerTask: