diff --git a/paasta_tools/cli/cmds/list_namespaces.py b/paasta_tools/cli/cmds/list_namespaces.py index fac36846cd..e5ed945925 100644 --- a/paasta_tools/cli/cmds/list_namespaces.py +++ b/paasta_tools/cli/cmds/list_namespaces.py @@ -15,6 +15,7 @@ from paasta_tools.cli.utils import get_instance_configs_for_service from paasta_tools.cli.utils import lazy_choices_completer from paasta_tools.cli.utils import validate_service_name +from paasta_tools.spark_tools import SPARK_EXECUTOR_NAMESPACE from paasta_tools.utils import DEFAULT_SOA_DIR from paasta_tools.utils import list_clusters from paasta_tools.utils import list_services @@ -69,8 +70,14 @@ def paasta_list_namespaces(args): # We skip non-k8s instance types if instance.get_instance_type() in ("paasta-native", "adhoc"): continue - else: - namespaces.add(instance.get_namespace()) + namespaces.add(instance.get_namespace()) + # Tron instances are TronActionConfigs + if ( + instance.get_instance_type() == "tron" + and instance.get_executor() == "spark" + ): + # We also need paasta-spark for spark executors + namespaces.add(SPARK_EXECUTOR_NAMESPACE) # Print in list format to be used in iam_roles print(list(namespaces)) diff --git a/tests/cli/test_cmds_list_namespaces.py b/tests/cli/test_cmds_list_namespaces.py index eb9bb635dc..eeb03910df 100644 --- a/tests/cli/test_cmds_list_namespaces.py +++ b/tests/cli/test_cmds_list_namespaces.py @@ -3,6 +3,7 @@ from mock import patch from paasta_tools.cli.cmds.list_namespaces import paasta_list_namespaces +from paasta_tools.spark_tools import SPARK_EXECUTOR_NAMESPACE def test_list_namespaces_no_instances(capfd): @@ -39,13 +40,14 @@ def create_mock_instance_config(instance_type, namespace): return mock_instance_config -def test_list_namespaces_with_instances(capfd): +def test_list_namespaces_with_instances_dupe_ns(capfd): mock_args = MagicMock( service="fake_service", instance=None, cluster=None, soa_dir="/fake/soa/dir", ) + mock_instance_configs = [ create_mock_instance_config("kubernetes", "k8s_namespace"), create_mock_instance_config("kubernetes", "k8s_namespace"), @@ -64,6 +66,59 @@ def test_list_namespaces_with_instances(capfd): assert stdout.strip() == "['k8s_namespace']" +def test_list_namespaces_tron(capfd): + mock_args = MagicMock( + service="fake_service", + instance=None, + cluster=None, + soa_dir="/fake/soa/dir", + ) + mock_tron_instance = create_mock_instance_config("tron", "tron") + mock_tron_instance.get_executor.return_value = "paasta" + + with patch( + "paasta_tools.cli.cmds.list_namespaces.get_instance_configs_for_service", + return_value=[mock_tron_instance], + autospec=True, + ), patch( + "paasta_tools.cli.cmds.list_namespaces.validate_service_name", + autospec=True, + ): + assert paasta_list_namespaces(mock_args) == 0 + stdout, _ = capfd.readouterr() + assert "['tron']" + + +def test_list_namespaces_spark(capfd): + mock_args = MagicMock( + service="fake_service", + instance=None, + cluster=None, + soa_dir="/fake/soa/dir", + ) + mock_spark_instance = create_mock_instance_config("tron", "tron") + mock_spark_instance.get_executor.return_value = "spark" + + mock_instance_configs = [ + create_mock_instance_config("kubernetes", "k8s_namespace"), + mock_spark_instance, + ] + + with patch( + "paasta_tools.cli.cmds.list_namespaces.get_instance_configs_for_service", + return_value=mock_instance_configs, + autospec=True, + ), patch( + "paasta_tools.cli.cmds.list_namespaces.validate_service_name", + autospec=True, + ): + assert paasta_list_namespaces(mock_args) == 0 + stdout, _ = capfd.readouterr() + assert f"'{SPARK_EXECUTOR_NAMESPACE}'" in stdout + assert "'tron'" in stdout + assert "'k8s_namespace'" in stdout + + def test_list_namespaces_skips_non_k8s_instances(capfd): mock_args = MagicMock( service="fake_service",