Fix pytest

NVIDIA · Nov 20, 2024 · 8e517e0 · 8e517e0
1 parent 9f581ea
commit 8e517e0
Showing 1 changed file with 16 additions and 5 deletions.
diff --git a/tests/slurm_command_gen_strategy/test_chakra_replay_slurm_command_gen_strategy.py b/tests/slurm_command_gen_strategy/test_chakra_replay_slurm_command_gen_strategy.py
@@ -31,12 +31,15 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> ChakraReplaySlurmComman
         return ChakraReplaySlurmCommandGenStrategy(slurm_system, {})
 
     @pytest.mark.parametrize(
-        "job_name_prefix, env_vars, cmd_args, num_nodes, nodes, expected_result",
+        "job_name_prefix, env_vars, test_definition_args, num_nodes, nodes, expected_result",
         [
             (
                 "chakra_replay",
                 {"NCCL_DEBUG": "INFO"},
-                {"docker_image_url": "fake_image_url", "trace_path": "/workspace/traces/"},
+                {
+                    "docker_image": Mock(installed_path="fake_image_url"),
+                    "cmd_args": Mock(trace_path="/workspace/traces/"),
+                },
                 2,
                 ["node1", "node2"],
                 {
@@ -47,7 +50,10 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> ChakraReplaySlurmComman
             (
                 "chakra_replay",
                 {"NCCL_DEBUG": "INFO"},
-                {"docker_image_url": "another_image_url", "trace_path": "/another/trace_path/"},
+                {
+                    "docker_image": Mock(installed_path="another_image_url"),
+                    "cmd_args": Mock(trace_path="/another/trace_path/"),
+                },
                 1,
                 ["node1"],
                 {
@@ -62,15 +68,20 @@ def test_parse_slurm_args(
         cmd_gen_strategy: ChakraReplaySlurmCommandGenStrategy,
         job_name_prefix: str,
         env_vars: Dict[str, str],
-        cmd_args: Dict[str, str],
+        test_definition_args: Dict[str, Any],
         num_nodes: int,
         nodes: List[str],
         expected_result: Dict[str, Any],
     ) -> None:
         tr = Mock(spec=TestRun)
         tr.num_nodes = num_nodes
         tr.nodes = nodes
-        slurm_args = cmd_gen_strategy._parse_slurm_args(job_name_prefix, env_vars, cmd_args, tr)
+        tr.test.test_definition = Mock(
+            docker_image=test_definition_args["docker_image"],
+            cmd_args=test_definition_args["cmd_args"],
+        )
+
+        slurm_args = cmd_gen_strategy._parse_slurm_args(job_name_prefix, env_vars, {}, tr)
         assert slurm_args["image_path"] == expected_result["image_path"]
         assert slurm_args["container_mounts"] == expected_result["container_mounts"]