Relocate MaskSecret comms tests into task SDK comms tests

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):

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"."""