TOOL-558: fix paasta status during migrations

paasta_tools/cli/cmds/status.py

@@ -2055,7 +2055,7 @@ def report_status_for_cluster(
     cluster: str,
     deploy_pipeline: Sequence[str],
     actual_deployments: Mapping[str, DeploymentVersion],
-    instance_whitelist: Mapping[str, Type[InstanceConfig]],
+    instance_whitelist: Mapping[str, List[Type[InstanceConfig]]],
     system_paasta_config: SystemPaastaConfig,
     lock: Lock,
     verbose: int = 0,
@@ -2068,16 +2068,17 @@ def report_status_for_cluster(
     deployed_instances = []
     instances = [
         (instance, instance_config_class)
-        for instance, instance_config_class in instance_whitelist.items()
+        for instance, config_classes in instance_whitelist.items()
+        for instance_config_class in config_classes
         if instance_config_class in ALLOWED_INSTANCE_CONFIG
     ]
 
     # Tron instance are not present in the deploy pipeline, so treat them as
     # seen by default to avoid error messages
     seen_instances = [
         instance
-        for instance, instance_config_class in instance_whitelist.items()
-        if instance_config_class == TronActionConfig
+        for instance, config_classes in instance_whitelist.items()
+        if TronActionConfig in config_classes
     ]
 
     for namespace in deploy_pipeline:
@@ -2094,9 +2095,9 @@ def report_status_for_cluster(
             deployed_instances.append(instance)
 
         # Case: flink/cassandra instances don't use `deployments.json`
-        elif (
-            instance_whitelist.get(instance)
-            in FLINK_DEPLOYMENT_CONFIGS + CASSANDRA_DEPLOYMENT_CONFIGS
+        elif instance in instance_whitelist and any(
+            cls in FLINK_DEPLOYMENT_CONFIGS + CASSANDRA_DEPLOYMENT_CONFIGS
+            for cls in instance_whitelist[instance]
         ):
             deployed_instances.append(instance)
 
@@ -2225,18 +2226,18 @@ def get_filters(
 
 def apply_args_filters(
     args,
-) -> Mapping[str, Mapping[str, Mapping[str, Type[InstanceConfig]]]]:
+) -> Mapping[str, Mapping[str, Mapping[str, List[Type[InstanceConfig]]]]]:
     """
     Take an args object and returns the dict of cluster:service:instances
     Currently, will filter by clusters, instances, services, and deploy_groups
     If no instances are found, will print a message and try to find matching instances
     for each service
 
     :param args: args object containing attributes to filter by
-    :returns: Dict of dicts, in format {cluster_name: {service_name: {instance1, instance2}}}
+    :returns: Dict of dicts, in format {cluster_name: {service_name: {instance1: [config_classes]}}}
     """
     clusters_services_instances: DefaultDict[
-        str, DefaultDict[str, Dict[str, Type[InstanceConfig]]]
+        str, DefaultDict[str, Dict[str, List[Type[InstanceConfig]]]]
     ] = defaultdict(lambda: defaultdict(dict))
     if args.service_instance:
         if args.service or args.instances:
@@ -2297,7 +2298,9 @@ def apply_args_filters(
                 cluster_service = clusters_services_instances[
                     instance_conf.get_cluster()
                 ][service]
-                cluster_service[instance_conf.get_instance()] = instance_conf.__class__
+                cluster_service.setdefault(instance_conf.get_instance(), []).append(
+                    instance_conf.__class__
+                )
                 i_count += 1
 
     if i_count == 0 and args.service and args.instances:
@@ -2319,9 +2322,14 @@ def paasta_status(args) -> int:
     clusters_services_instances = apply_args_filters(args)
     for cluster, service_instances in clusters_services_instances.items():
         for service, instances in service_instances.items():
-            all_flink = all((i in FLINK_DEPLOYMENT_CONFIGS) for i in instances.values())
+            all_config_classes = [
+                cls for classes in instances.values() for cls in classes
+            ]
+            all_flink = all(
+                cls in FLINK_DEPLOYMENT_CONFIGS for cls in all_config_classes
+            )
             all_cassandra = all(
-                (i in CASSANDRA_DEPLOYMENT_CONFIGS) for i in instances.values()
+                cls in CASSANDRA_DEPLOYMENT_CONFIGS for cls in all_config_classes
             )
             actual_deployments: Mapping[str, DeploymentVersion]
             if all_flink or all_cassandra:

tests/cli/test_cmds_status.py

@@ -28,6 +28,8 @@
 import paasta_tools.paastaapi.models as paastamodels
 from paasta_tools import kubernetes_tools
 from paasta_tools import utils
+from paasta_tools.cassandracluster_tools import CassandraClusterDeploymentConfig
+from paasta_tools.cassandraclustereks_tools import CassandraClusterEksDeploymentConfig
 from paasta_tools.cli.cmds import status
 from paasta_tools.cli.cmds.status import OUTPUT_HORIZONTAL_RULE
 from paasta_tools.cli.cmds.status import append_pod_status
@@ -352,7 +354,7 @@ def test_status_calls_sergeants(
         deploy_pipeline=planned_deployments,
         actual_deployments=actual_deployments,
         cluster=cluster,
-        instance_whitelist={"fi": mock_instance_config.__class__},
+        instance_whitelist={"fi": [mock_instance_config.__class__]},
         system_paasta_config=system_paasta_config,
         lock=mock.ANY,
         verbose=False,
@@ -453,7 +455,9 @@ def test_apply_args_filters_clusters_and_instances_clusters_instances_deploy_gro
 
     pargs = apply_args_filters(args)
     assert sorted(pargs.keys()) == ["cluster1"]
-    assert pargs["cluster1"]["fake_service"] == {"instance1": mock_inst1.__class__}
+    assert pargs["cluster1"]["fake_service"] == {
+        "instance1": [mock_inst1.__class__],
+    }
 
 
 @patch("paasta_tools.cli.cmds.status.get_instance_configs_for_service", autospec=True)
@@ -501,10 +505,12 @@ def test_apply_args_filters_clusters_uses_deploy_group_when_no_clusters_and_inst
     pargs = apply_args_filters(args)
     assert sorted(pargs.keys()) == ["cluster1", "cluster2"]
     assert pargs["cluster1"]["fake_service"] == {
-        "instance1": mock_inst1.__class__,
-        "instance2": mock_inst2.__class__,
+        "instance1": [mock_inst1.__class__],
+        "instance2": [mock_inst2.__class__],
+    }
+    assert pargs["cluster2"]["fake_service"] == {
+        "instance3": [mock_inst3.__class__],
     }
-    assert pargs["cluster2"]["fake_service"] == {"instance3": mock_inst3.__class__}
 
 
 @patch("paasta_tools.cli.cmds.status.get_instance_configs_for_service", autospec=True)
@@ -627,8 +633,8 @@ def test_apply_args_filters_clusters_and_instances(
     pargs = apply_args_filters(args)
     assert sorted(pargs.keys()) == ["cluster1"]
     assert pargs["cluster1"]["fake_service"] == {
-        "instance1": mock_inst1.__class__,
-        "instance3": mock_inst3.__class__,
+        "instance1": [mock_inst1.__class__],
+        "instance3": [mock_inst3.__class__],
     }
 
 
@@ -679,12 +685,14 @@ def test_apply_args_filters_shorthand_notation(
     pargs = apply_args_filters(args)
     if service_instance_name == "fake_service.instance1":
         assert sorted(pargs.keys()) == ["cluster1"]
-        assert pargs["cluster1"]["fake_service"] == {"instance1": mock_inst1.__class__}
+        assert pargs["cluster1"]["fake_service"] == {
+            "instance1": [mock_inst1.__class__],
+        }
     elif service_instance_name == "fake_service.instance1,instance2":
         assert sorted(pargs.keys()) == ["cluster1"]
         assert pargs["cluster1"]["fake_service"] == {
-            "instance1": mock_inst1.__class__,
-            "instance2": mock_inst2.__class__,
+            "instance1": [mock_inst1.__class__],
+            "instance2": [mock_inst2.__class__],
         }
     elif service_instance_name == "fake_service.instance3":
         assert sorted(pargs.keys()) == []
@@ -769,6 +777,57 @@ def test_apply_args_filters_no_instances_found(
         assert i in output
 
 
+@patch("paasta_tools.cli.cmds.status.get_instance_configs_for_service", autospec=True)
+@patch("paasta_tools.cli.cmds.status.list_services", autospec=True)
+@patch("paasta_tools.cli.cmds.status.figure_out_service_name", autospec=True)
+@patch("paasta_tools.cli.cmds.status.validate_service_name", autospec=True)
+def test_apply_args_filters_returns_both_cassandra_and_cassandraeks(
+    mock_validate_service_name,
+    mock_figure_out_service_name,
+    mock_list_services,
+    mock_get_instance_configs_for_service,
+):
+    """When both cassandracluster and cassandraclustereks configs exist for the
+    same instance name, apply_args_filters should return both config classes."""
+    args = StatusArgs(
+        service="fake_service",
+        soa_dir="/fake/soa/dir",
+        deploy_group=None,
+        clusters="cluster1",
+        instances="instance1",
+        owner=None,
+        registration=None,
+        verbose=False,
+        service_instance=None,
+    )
+    mock_validate_service_name.return_value = None
+    mock_figure_out_service_name.return_value = "fake_service"
+    mock_list_services.return_value = ["fake_service"]
+
+    mock_cassandra_inst = make_fake_instance_conf(
+        "cluster1", "fake_service", "instance1"
+    )
+    mock_cassandra_inst.__class__ = CassandraClusterDeploymentConfig
+    mock_cassandraeks_inst = make_fake_instance_conf(
+        "cluster1", "fake_service", "instance1"
+    )
+    mock_cassandraeks_inst.__class__ = CassandraClusterEksDeploymentConfig
+
+    mock_get_instance_configs_for_service.return_value = [
+        mock_cassandra_inst,
+        mock_cassandraeks_inst,
+    ]
+
+    pargs = apply_args_filters(args)
+    assert sorted(pargs.keys()) == ["cluster1"]
+    assert pargs["cluster1"]["fake_service"] == {
+        "instance1": [
+            CassandraClusterDeploymentConfig,
+            CassandraClusterEksDeploymentConfig,
+        ],
+    }
+
+
 @patch("paasta_tools.cli.cmds.status.get_instance_configs_for_service", autospec=True)
 @patch("paasta_tools.cli.cmds.status.list_services", autospec=True)
 @patch("paasta_tools.cli.cmds.status.figure_out_service_name", autospec=True)
@@ -907,8 +966,8 @@ def test_status_with_registration(
         deploy_pipeline=ANY,
         actual_deployments=ANY,
         instance_whitelist={
-            "instance1": mock_inst_1.__class__,
-            "instance2": mock_inst_2.__class__,
+            "instance1": [mock_inst_1.__class__],
+            "instance2": [mock_inst_2.__class__],
         },
         system_paasta_config=system_paasta_config,
         lock=mock.ANY,