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 6d73c623f5ea6..c76fa76d90f79 100644 --- a/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py +++ b/task-sdk/tests/task_sdk/definitions/test_secrets_masker.py @@ -30,7 +30,6 @@ import pytest from airflow.models import Connection -from airflow.sdk.execution_time.comms import MaskSecret from airflow.sdk.execution_time.secrets_masker import ( RedactedIO, SecretsMasker, @@ -565,36 +564,6 @@ def test_add_mask_short_secrets_and_skip_keywords( if should_be_masked: assert filt.replacer is not None - @pytest.mark.parametrize( - "object_to_mask", - [ - { - "key_path": "/files/airflow-breeze-config/keys2/keys.json", - "scope": "https://www.googleapis.com/auth/cloud-platform", - "project": "project_id", - "num_retries": 6, - }, - ["iter1", "iter2", {"key": "value"}], - "string", - { - "key1": "value1", - }, - ], - ) - def test_mask_secret_with_objects(self, object_to_mask): - mask_secret_object = MaskSecret(value=object_to_mask, name="test_secret") - assert mask_secret_object.value == object_to_mask - - def test_mask_secret_with_list(self): - example_dict = ["test"] - mask_secret_object = MaskSecret(value=example_dict, name="test_secret") - assert mask_secret_object.value == example_dict - - def test_mask_secret_with_iterable(self): - example_dict = ["test"] - mask_secret_object = MaskSecret(value=example_dict, name="test_secret") - assert mask_secret_object.value == example_dict - class TestStructuredVsUnstructuredMasking: def test_structured_sensitive_fields_always_masked(self): diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py b/task-sdk/tests/task_sdk/execution_time/test_comms.py index 2f5d6bceedc21..fd5d352e14ddf 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -25,10 +25,44 @@ import pytest from airflow.sdk import timezone -from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails, _ResponseFrame +from airflow.sdk.execution_time.comms import BundleInfo, MaskSecret, StartupDetails, _ResponseFrame from airflow.sdk.execution_time.task_runner import CommsDecoder +class TestCommsModels: + """Test Pydantic models used in task communication for proper validation.""" + + @pytest.mark.parametrize( + "object_to_mask", + [ + { + "key_path": "/files/airflow-breeze-config/keys2/keys.json", + "scope": "https://www.googleapis.com/auth/cloud-platform", + "project": "project_id", + "num_retries": 6, + }, + ["iter1", "iter2", {"key": "value"}], + "string", + { + "key1": "value1", + }, + ], + ) + def test_mask_secret_with_objects(self, object_to_mask): + mask_secret_object = MaskSecret(value=object_to_mask, name="test_secret") + assert mask_secret_object.value == object_to_mask + + def test_mask_secret_with_list(self): + example_dict = ["test"] + mask_secret_object = MaskSecret(value=example_dict, name="test_secret") + assert mask_secret_object.value == example_dict + + def test_mask_secret_with_iterable(self): + example_dict = ["test"] + mask_secret_object = MaskSecret(value=example_dict, name="test_secret") + assert mask_secret_object.value == example_dict + + class TestCommsDecoder: """Test the communication between the subprocess and the "supervisor"."""