Improve Argo workflow label tests

Author: talsperre

test/unit/test_argo_workflows_labels.py

@@ -1,246 +1,64 @@
 import pytest
-from unittest.mock import patch, MagicMock, PropertyMock
 
+from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
 from metaflow.plugins.kubernetes.kube_utils import KubernetesException
 
 
-class TestArgoWorkflowsLabels:
-    """Test _base_argo_labels() method with mocked configuration."""
+@pytest.fixture
+def argo_workflows():
+    return ArgoWorkflows.__new__(ArgoWorkflows)
 
-    def _call_base_argo_labels(self):
-        """
-        Call _base_argo_labels without full ArgoWorkflows instantiation.
 
-        The methods only use module-level config, not instance state, so we can
-        call them on a minimal object with the methods bound to it.
-        """
-        from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
-
-        # Create a minimal object with both methods bound
-        class MinimalArgo:
-            _base_kubernetes_labels = ArgoWorkflows._base_kubernetes_labels
-            _custom_argo_labels = ArgoWorkflows._custom_argo_labels
-            _base_argo_labels = ArgoWorkflows._base_argo_labels
-
-        obj = MinimalArgo()
-        return obj._base_argo_labels()
-
-    def test_default_labels_when_env_not_set(self):
-        """Should return default labels when ARGO_WORKFLOWS_LABELS is empty."""
-        with patch("metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS", ""):
-            labels = self._call_base_argo_labels()
-
-        assert labels == {"app.kubernetes.io/part-of": "metaflow"}
-
-    def test_adds_custom_labels_from_env(self):
-        """Should add custom labels from ARGO_WORKFLOWS_LABELS."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
+@pytest.mark.parametrize(
+    ("configured_labels", "expected"),
+    [
+        (
+            "",
+            {"app.kubernetes.io/part-of": "metaflow"},
+        ),
+        (
             "team=ml,env=prod",
-        ):
-            labels = self._call_base_argo_labels()
-
-        assert labels == {
-            "app.kubernetes.io/part-of": "metaflow",
-            "team": "ml",
-            "env": "prod",
-        }
-
-    def test_custom_labels_do_not_override_internal_labels(self):
-        """Custom labels must not override internal/base labels."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            "app.kubernetes.io/part-of=custom-app",
-        ):
-            labels = self._call_base_argo_labels()
-
-        assert labels == {"app.kubernetes.io/part-of": "metaflow"}
-
-    def test_custom_labels_alongside_part_of_override_attempt(self):
-        """Unrelated custom labels pass through even when part-of override is attempted."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            "app.kubernetes.io/part-of=custom-app,team=ml",
-        ):
-            labels = self._call_base_argo_labels()
-
-        assert labels == {
-            "app.kubernetes.io/part-of": "metaflow",
-            "team": "ml",
-        }
-
-    def test_single_label(self):
-        """Should handle a single label correctly."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            "cost-center=12345",
-        ):
-            labels = self._call_base_argo_labels()
-
-        assert labels == {
-            "app.kubernetes.io/part-of": "metaflow",
-            "cost-center": "12345",
-        }
-
-    def test_invalid_label_value_raises_exception(self):
-        """Should raise exception for invalid label values."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            "team=invalid value with spaces",
-        ):
-            with pytest.raises(KubernetesException):
-                self._call_base_argo_labels()
-
-    def test_label_value_too_long_raises_exception(self):
-        """Should raise exception for label values exceeding 63 chars."""
-        long_value = "a" * 64
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            f"team={long_value}",
-        ):
-            with pytest.raises(KubernetesException):
-                self._call_base_argo_labels()
-
-    def test_invalid_label_key_empty_name_raises_exception(self):
-        """Should raise exception for a key with an empty name segment (e.g. '=ml')."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            "=ml",
-        ):
-            with pytest.raises(KubernetesException):
-                self._call_base_argo_labels()
-
-    def test_invalid_label_key_with_spaces_raises_exception(self):
-        """Should raise exception for a key containing spaces."""
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            " team=ml",
-        ):
-            with pytest.raises(KubernetesException):
-                self._call_base_argo_labels()
-
-    def test_invalid_label_key_too_long_raises_exception(self):
-        """Should raise exception for a key name segment exceeding 63 chars."""
-        long_key = "a" * 64
-        with patch(
-            "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-            f"{long_key}=ml",
-        ):
-            with pytest.raises(KubernetesException):
-                self._call_base_argo_labels()
-
-
-class TestArgoWorkflowsTemplateLabels:
-    """Test that labels appear in compiled WorkflowTemplate and Sensor."""
-
-    @pytest.fixture
-    def mock_argo_workflows(self):
-        """Create an ArgoWorkflows instance with mocked dependencies."""
-        # Mock all the complex dependencies
-        patches = [
-            patch(
-                "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
-                "team=ml-platform,env=production",
-            ),
-            patch(
-                "metaflow.plugins.argo.argo_workflows.KUBERNETES_NAMESPACE", "test-ns"
-            ),
-            patch("metaflow.plugins.argo.argo_workflows.ARGO_EVENTS_EVENT", None),
-            patch(
-                "metaflow.plugins.argo.argo_workflows.ARGO_EVENTS_EVENT_SOURCE", None
-            ),
-            patch(
-                "metaflow.plugins.argo.argo_workflows.ARGO_EVENTS_SERVICE_ACCOUNT", None
-            ),
-        ]
-
-        for p in patches:
-            p.start()
-
-        from metaflow.plugins.argo.argo_workflows import ArgoWorkflows
-
-        # Create mock graph with minimal structure
-        mock_node = MagicMock()
-        mock_node.name = "start"
-        mock_node.type = "linear"
-        mock_node.out_funcs = ["end"]
-        mock_node.is_inside_foreach = False
-        mock_node.parallel_foreach = False
-
-        mock_end_node = MagicMock()
-        mock_end_node.name = "end"
-        mock_end_node.type = "end"
-        mock_end_node.out_funcs = []
-        mock_end_node.is_inside_foreach = False
-        mock_end_node.parallel_foreach = False
-
-        mock_graph = MagicMock()
-        mock_graph.nodes = {"start": mock_node, "end": mock_end_node}
-        mock_graph.__iter__ = lambda self: iter([mock_node, mock_end_node])
-
-        # Create mock flow
-        mock_flow = MagicMock()
-        mock_flow.name = "TestFlow"
-        mock_flow._flow_decorators = {}
-        type(mock_flow)._parameters = PropertyMock(return_value={})
-        type(mock_flow)._configs = PropertyMock(return_value={})
-
-        # Create mock environment
-        mock_environment = MagicMock()
-        mock_environment.get_package_commands.return_value = []
-        mock_environment.bootstrap_commands.return_value = []
-
-        # Create mock datastore
-        mock_datastore = MagicMock()
-        mock_datastore.TYPE = "s3"
-
-        # Create the instance - we'll patch the complex compilation methods
-        with patch.object(ArgoWorkflows, "_compile_workflow_template"), patch.object(
-            ArgoWorkflows, "_compile_sensor"
-        ), patch.object(
-            ArgoWorkflows, "_process_parameters", return_value=[]
-        ), patch.object(
-            ArgoWorkflows, "_process_config_parameters", return_value=[]
-        ), patch.object(
-            ArgoWorkflows, "_process_triggers", return_value=([], {})
-        ), patch.object(
-            ArgoWorkflows, "_get_schedule", return_value=(None, None)
-        ), patch.object(
-            ArgoWorkflows, "_parse_conditional_branches"
-        ):
-            argo = ArgoWorkflows(
-                name="test-flow",
-                graph=mock_graph,
-                flow=mock_flow,
-                code_package_metadata={},
-                code_package_sha="abc123",
-                code_package_url="s3://bucket/code.tar.gz",
-                production_token="prod-token",
-                metadata=MagicMock(),
-                flow_datastore=mock_datastore,
-                environment=mock_environment,
-                event_logger=MagicMock(),
-                monitor=MagicMock(),
-                username="testuser",
-            )
-
-        for p in patches:
-            p.stop()
-
-        return argo
-
-    def test_workflow_labels_includes_custom_labels(self, mock_argo_workflows):
-        """Verify _workflow_labels (used at WorkflowTemplate/Workflow level) contains custom labels from env var."""
-        labels = mock_argo_workflows._workflow_labels
-
-        assert labels["app.kubernetes.io/part-of"] == "metaflow"
-        assert labels["team"] == "ml-platform"
-        assert labels["env"] == "production"
-
-    def test_base_labels_excludes_custom_labels(self, mock_argo_workflows):
-        """Verify _base_labels (used for pods/JobSet/Sensor) does NOT contain custom labels from env var."""
-        labels = mock_argo_workflows._base_labels
-
-        assert labels == {"app.kubernetes.io/part-of": "metaflow"}
-        assert "team" not in labels
-        assert "env" not in labels
+            {
+                "app.kubernetes.io/part-of": "metaflow",
+                "team": "ml",
+                "env": "prod",
+            },
+        ),
+        (
+            "app.kubernetes.io/part-of=custom,team=ml",
+            {
+                "app.kubernetes.io/part-of": "metaflow",
+                "team": "ml",
+            },
+        ),
+    ],
+    ids=["default", "custom-labels", "protected-label"],
+)
+def test_base_argo_labels(mocker, argo_workflows, configured_labels, expected):
+    mocker.patch(
+        "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
+        configured_labels,
+    )
+
+    assert argo_workflows._base_argo_labels() == expected
+
+
+@pytest.mark.parametrize(
+    "configured_labels",
+    [
+        "missing-value",
+        "team=value with spaces",
+        "team=%s" % ("a" * 64),
+    ],
+    ids=["missing-equals", "invalid-value", "value-too-long"],
+)
+def test_base_argo_labels_rejects_invalid_configuration(
+    mocker, argo_workflows, configured_labels
+):
+    mocker.patch(
+        "metaflow.plugins.argo.argo_workflows.ARGO_WORKFLOWS_LABELS",
+        configured_labels,
+    )
+
+    with pytest.raises(KubernetesException):
+        argo_workflows._base_argo_labels()

test/ux/core/test_argo_compilation.py

@@ -64,6 +64,55 @@ def test_argo_only_json_exposes_workflow_template(
     assert workflow_template["spec"]["templates"]
 
 
+def test_configured_labels_are_emitted_on_argo_workflows(
+    exec_mode, decospecs, compute_env, tag, scheduler_config
+):
+    if exec_mode != "deployer":
+        pytest.skip("Argo compilation tests require deployer mode")
+    if scheduler_config.scheduler_type != "argo-workflows":
+        pytest.skip("Argo compilation tests require the argo-workflows scheduler")
+
+    from metaflow import Deployer
+
+    from .test_utils import _resolve_flow_path, prepare_runner_deployer_args
+
+    env = dict(compute_env)
+    env["METAFLOW_ARGO_WORKFLOWS_LABELS"] = (
+        "team=ml-platform,environment=test,app.kubernetes.io/name=custom"
+    )
+
+    deployed_flow = (
+        Deployer(
+            flow_file=_resolve_flow_path("basic/helloworld.py"),
+            show_output=False,
+            **prepare_runner_deployer_args({"decospecs": decospecs, "env": env}),
+        )
+        .argo_workflows()
+        .create(
+            only_json=True,
+            tags=tag + ["test_configured_argo_labels"],
+            **(scheduler_config.deploy_args or {}),
+        )
+    )
+
+    workflow_template = deployed_flow.workflow_template
+    template_labels = workflow_template["metadata"]["labels"]
+    workflow_labels = workflow_template["spec"]["workflowMetadata"]["labels"]
+
+    assert template_labels["team"] == "ml-platform"
+    assert template_labels["environment"] == "test"
+    assert template_labels["app.kubernetes.io/name"] == "metaflow-flow"
+
+    assert workflow_labels["team"] == "ml-platform"
+    assert workflow_labels["environment"] == "test"
+    assert workflow_labels["app.kubernetes.io/name"] == "metaflow-run"
+
+    # Custom Argo labels are intentionally workflow-level only.
+    pod_labels = workflow_template["spec"]["podMetadata"]["labels"]
+    assert "team" not in pod_labels
+    assert "environment" not in pod_labels
+
+
 def test_foreach_split_switch_join_task_names_are_deduplicated(
     exec_mode, decospecs, tag, scheduler_config
 ):