Fix pytest

NVIDIA · Nov 20, 2024 · b18d109 · b18d109
1 parent 8e517e0
commit b18d109
Showing 1 changed file with 8 additions and 14 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
@@ -22,6 +22,7 @@
 from cloudai import TestRun
 from cloudai.schema.test_template.chakra_replay.slurm_command_gen_strategy import ChakraReplaySlurmCommandGenStrategy
 from cloudai.systems import SlurmSystem
+from cloudai.test_definitions.chakra_replay import ChakraReplayCmdArgs, ChakraReplayTestDefinition
 from tests.conftest import create_autospec_dataclass
 
 
@@ -31,15 +32,12 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> ChakraReplaySlurmComman
         return ChakraReplaySlurmCommandGenStrategy(slurm_system, {})
 
     @pytest.mark.parametrize(
-        "job_name_prefix, env_vars, test_definition_args, num_nodes, nodes, expected_result",
+        "job_name_prefix, env_vars, cmd_args_attrs, num_nodes, nodes, expected_result",
         [
             (
                 "chakra_replay",
                 {"NCCL_DEBUG": "INFO"},
-                {
-                    "docker_image": Mock(installed_path="fake_image_url"),
-                    "cmd_args": Mock(trace_path="/workspace/traces/"),
-                },
+                {"docker_image_url": "fake_image_url", "trace_path": "/workspace/traces/"},
                 2,
                 ["node1", "node2"],
                 {
@@ -50,10 +48,7 @@ def cmd_gen_strategy(self, slurm_system: SlurmSystem) -> ChakraReplaySlurmComman
             (
                 "chakra_replay",
                 {"NCCL_DEBUG": "INFO"},
-                {
-                    "docker_image": Mock(installed_path="another_image_url"),
-                    "cmd_args": Mock(trace_path="/another/trace_path/"),
-                },
+                {"docker_image_url": "another_image_url", "trace_path": "/another/trace_path/"},
                 1,
                 ["node1"],
                 {
@@ -68,18 +63,17 @@ def test_parse_slurm_args(
         cmd_gen_strategy: ChakraReplaySlurmCommandGenStrategy,
         job_name_prefix: str,
         env_vars: Dict[str, str],
-        test_definition_args: Dict[str, Any],
+        cmd_args_attrs: Dict[str, str],
         num_nodes: int,
         nodes: List[str],
         expected_result: Dict[str, Any],
     ) -> None:
+        cmd_args = ChakraReplayCmdArgs(**cmd_args_attrs)
+        test_definition = ChakraReplayTestDefinition(cmd_args=cmd_args)
         tr = Mock(spec=TestRun)
         tr.num_nodes = num_nodes
         tr.nodes = nodes
-        tr.test.test_definition = Mock(
-            docker_image=test_definition_args["docker_image"],
-            cmd_args=test_definition_args["cmd_args"],
-        )
+        tr.test.test_definition = test_definition
 
         slurm_args = cmd_gen_strategy._parse_slurm_args(job_name_prefix, env_vars, {}, tr)
         assert slurm_args["image_path"] == expected_result["image_path"]