Fix unit tests

Author: hgreebe

src/slurm_plugin/instance_manager.py

@@ -848,7 +848,11 @@ def _best_effort_node_assignment(
             # No instances launched at all, e.g. CreateFleet API returns no EC2 instances,
             # or no left instances available from a best-effort EC2 launch
             logger.info("No launched instances found for nodes %s", print_with_count(nodes_resume_list))
-            self._update_failed_nodes(set(nodes_resume_list), "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)", override=False)
+            self._update_failed_nodes(
+                set(nodes_resume_list),
+                "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)",
+                override=False,
+            )
 
     def _all_or_nothing_node_assignment(
         self,
@@ -903,7 +907,11 @@ def _all_or_nothing_node_assignment(
             # No instances launched at all, e.g. CreateFleet API returns no EC2 instances,
             # or no left instances available from a best-effort EC2 launch
             logger.info("No launched instances found for nodes %s", print_with_count(nodes_resume_list))
-            self._update_failed_nodes(set(nodes_resume_list), "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)", override=False)
+            self._update_failed_nodes(
+                set(nodes_resume_list),
+                "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)",
+                override=False,
+            )
 
     def _launch_instances(  # noqa: C901
         self,
@@ -986,7 +994,9 @@ def _launch_ec2_instances(self, batch_nodes, compute_resource, fleet_manager, in
             #   queue_2: {cr_3: list[EC2Instance]}
             # }
         else:
-            self._update_failed_nodes(set(batch_nodes), "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)")
+            self._update_failed_nodes(
+                set(batch_nodes), "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)"
+            )
         return launched_ec2_instances
 
     def _get_fleet_manager(self, all_or_nothing_batch, compute_resource, queue):

tests/slurm_plugin/test_instance_manager.py

@@ -3195,7 +3195,7 @@ def test_update_slurm_node_addrs(
                 {},
                 False,
                 None,
-                {"InsufficientInstanceCapacity": {"queue4-st-c5xlarge-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"queue4-st-c5xlarge-1"}},
             ),
             (
                 SlurmResumeJob(140819, "queue4-st-c5xlarge-1", "queue4-st-c5xlarge-1", "NO"),
@@ -3210,7 +3210,7 @@ def test_update_slurm_node_addrs(
                 {},
                 False,
                 None,
-                {"InsufficientInstanceCapacity": {"queue4-st-c5xlarge-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"queue4-st-c5xlarge-1"}},
             ),
             (
                 SlurmResumeJob(140819, "queue4-st-c5xlarge-1", "queue4-st-c5xlarge-1", "NO"),
@@ -3328,7 +3328,12 @@ def test_update_slurm_node_addrs(
                 {},
                 False,
                 None,
-                {"InsufficientInstanceCapacity": {"queue1-st-c5xlarge-1", "queue4-st-c5xlarge-1"}},
+                {
+                    "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {
+                        "queue1-st-c5xlarge-1",
+                        "queue4-st-c5xlarge-1",
+                    }
+                },
             ),
             (
                 SlurmResumeJob(
@@ -3348,7 +3353,7 @@ def test_update_slurm_node_addrs(
                 {},
                 False,
                 None,
-                {"InsufficientInstanceCapacity": {"queue1-st-c5xlarge-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"queue1-st-c5xlarge-1"}},
             ),
             (
                 SlurmResumeJob(
@@ -3915,7 +3920,7 @@ def test_node_assignment_by_scaling_strategy(
                 {},
                 [{"Instances": []}],
                 {},
-                {"InsufficientInstanceCapacity": {"queue4-st-c5xlarge-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"queue4-st-c5xlarge-1"}},
             ),
             (
                 None,
@@ -4719,7 +4724,7 @@ def test_update_dict(self, instance_manager, target_dict, update, expected_dict)
                 {"q1": {"c1": ["q1-dy-c1-1"], "c2": ["q1-dy-c2-1"]}},
                 [],
                 True,
-                {"InsufficientInstanceCapacity": {"q1-dy-c1-1", "q1-dy-c2-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"q1-dy-c1-1", "q1-dy-c2-1"}},
             ),
         ],
     )
@@ -4898,7 +4903,7 @@ def test_best_effort_node_assignment(
                 {"q1": {"c1": ["q1-dy-c1-1"], "c2": ["q1-dy-c2-1"]}},
                 [],
                 True,
-                {"InsufficientInstanceCapacity": {"q1-dy-c1-1", "q1-dy-c2-1"}},
+                {"InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {"q1-dy-c1-1", "q1-dy-c2-1"}},
                 None,
                 {},
                 {},
@@ -5590,7 +5595,7 @@ def test_add_instances(
                 {},
                 False,
                 {
-                    "InsufficientInstanceCapacity": {
+                    "InsufficientInstanceCapacity(Check slurm_resume log for ec2 error codes)": {
                         "queue2-st-c5xlarge-1",
                         "queue2-dy-c5xlarge-1",
                         "queue1-st-c5xlarge-2",