Add unit test to cover max_count > 0 in _is_node_in_replacement_valid

Author: dreambeyondorange

tests/slurm_plugin/test_clustermgtd.py

@@ -2648,43 +2648,61 @@ def initialize_console_logger_mock(mocker):
 
 
 @pytest.mark.parametrize(
-    "current_replacing_nodes, node, instance, current_time, expected_result",
+    "current_replacing_nodes, node, instance, current_time, max_count, expected_result",
     [
         (
             set(),
             StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
             EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
             datetime(2020, 1, 1, 0, 0, 29),
+            0,
             False,
         ),
         (
             {"queue1-st-c5xlarge-1"},
             StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
             None,
             datetime(2020, 1, 1, 0, 0, 29),
+            0,
             False,
         ),
         (
             {"queue1-st-c5xlarge-1"},
             StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD", "queue1"),
             EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
             datetime(2020, 1, 1, 0, 0, 29),
+            0,
             True,
         ),
         (
             {"queue1-st-c5xlarge-1"},
             StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
             EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
             datetime(2020, 1, 1, 0, 0, 30),
+            0,
             False,
         ),
+        (
+            {"queue1-st-c5xlarge-1"},
+            StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
+            None,
+            datetime(2020, 1, 1, 0, 0, 30),
+            1,
+            True,
+        ),
+    ],
+    ids=[
+        "not_in_replacement",
+        "no-backing-instance",
+        "in_replacement",
+        "timeout",
+        "no-backing-instance-with-max-count",
     ],
-    ids=["not_in_replacement", "no-backing-instance", "in_replacement", "timeout"],
 )
 @pytest.mark.usefixtures(
     "initialize_instance_manager_mock", "initialize_executor_mock", "initialize_console_logger_mock"
 )
-def test_is_node_being_replaced(current_replacing_nodes, node, instance, current_time, expected_result):
+def test_is_node_being_replaced(current_replacing_nodes, node, instance, current_time, max_count, expected_result):
     mock_sync_config = SimpleNamespace(
         node_replacement_timeout=30,
         insufficient_capacity_timeout=3,
@@ -2693,7 +2711,7 @@ def test_is_node_being_replaced(current_replacing_nodes, node, instance, current
         region="region",
         boto3_config=None,
         fleet_config={},
-        ec2_instance_missing_max_count=0,
+        ec2_instance_missing_max_count=max_count,
     )
     cluster_manager = ClusterManager(mock_sync_config)
     cluster_manager._current_time = current_time