diff --git a/task-sdk/tests/conftest.py b/task-sdk/tests/conftest.py index f668ed911a21f..9e9efd6c928e4 100644 --- a/task-sdk/tests/conftest.py +++ b/task-sdk/tests/conftest.py @@ -20,6 +20,7 @@ import os from pathlib import Path from typing import TYPE_CHECKING, Any, NoReturn, Protocol +from unittest.mock import patch import pytest @@ -266,3 +267,12 @@ def _make_context_dict( return context.model_dump(exclude_unset=True, mode="json") return _make_context_dict + + +@pytest.fixture(autouse=True) +def set_secrets_masker(): + from airflow.sdk.execution_time.secrets_masker import SecretsMasker + + secrets_masker = SecretsMasker() + with patch("airflow.sdk.execution_time.secrets_masker._secrets_masker", return_value=secrets_masker): + yield secrets_masker diff --git a/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py b/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py index 7ad1f4a70f336..f63daa51a19f2 100644 --- a/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py +++ b/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py @@ -432,14 +432,10 @@ def lineno(): class TestRedactedIO: - @pytest.fixture(scope="class", autouse=True) - def reset_secrets_masker(self): - self.secrets_masker = SecretsMasker() - with patch( - "airflow.sdk.execution_time.secrets_masker._secrets_masker", return_value=self.secrets_masker - ): - mask_secret(p) - yield + @pytest.fixture(autouse=True) + def reset_secrets_masker(self, set_secrets_masker): + mask_secret(p) + return def test_redacts_from_print(self, capsys): # Without redacting, password is printed. @@ -472,33 +468,29 @@ def test_input_builtin(self, monkeypatch): class TestMaskSecretAdapter: @pytest.fixture(autouse=True) - def reset_secrets_masker_and_skip_escape(self): - self.secrets_masker = SecretsMasker() - with patch( - "airflow.sdk.execution_time.secrets_masker._secrets_masker", return_value=self.secrets_masker - ): - with patch("airflow.sdk.execution_time.secrets_masker.re.escape", lambda x: x): - yield + def patched_escape(self, set_secrets_masker): + with patch("airflow.sdk.execution_time.secrets_masker.re.escape", lambda x: x): + yield set_secrets_masker - def test_calling_mask_secret_adds_adaptations_for_returned_str(self): + def test_calling_mask_secret_adds_adaptations_for_returned_str(self, patched_escape): with conf_vars({("logging", "secret_mask_adapter"): "urllib.parse.quote"}): mask_secret("secret<>&", None) - assert self.secrets_masker.patterns == {"secret%3C%3E%26", "secret<>&"} + assert patched_escape.patterns == {"secret%3C%3E%26", "secret<>&"} - def test_calling_mask_secret_adds_adaptations_for_returned_iterable(self): + def test_calling_mask_secret_adds_adaptations_for_returned_iterable(self, patched_escape): with conf_vars({("logging", "secret_mask_adapter"): "urllib.parse.urlparse"}): mask_secret("https://airflow.apache.org/docs/apache-airflow/stable", "password") - assert self.secrets_masker.patterns == { + assert patched_escape.patterns == { "https", "airflow.apache.org", "/docs/apache-airflow/stable", "https://airflow.apache.org/docs/apache-airflow/stable", } - def test_calling_mask_secret_not_set(self): + def test_calling_mask_secret_not_set(self, patched_escape): with conf_vars({("logging", "secret_mask_adapter"): None}): mask_secret("a secret") - assert self.secrets_masker.patterns == {"a secret"} + assert patched_escape.patterns == {"a secret"}