Rename code_command.py to code.py, fix test

Signed-off-by: Fabrice Normandin <[email protected]>

milatools/cli/commands.py

@@ -30,7 +30,7 @@
 from ..__version__ import __version__
 from ..utils.local_v1 import LocalV1
 from ..utils.remote_v1 import RemoteV1
-from .code_command import add_mila_code_arguments
+from .code import add_mila_code_arguments
 from .common import forward, standard_server
 from .init_command import (
     print_welcome_message,

tests/integration/test_code_command.py → tests/integration/test_code.py

@@ -13,7 +13,7 @@
 import pytest
 from pytest_regressions.file_regression import FileRegressionFixture
 
-from milatools.cli.code_command import code
+from milatools.cli.code import code
 from milatools.cli.common import check_disk_quota
 from milatools.cli.utils import get_hostname_to_use_for_compute_node
 from milatools.utils.compute_node import ComputeNode
@@ -103,24 +103,24 @@ async def get_job_info(
     indirect=True,
 )
 async def test_code(
-    login_node: RemoteV2,
+    login_node_v2: RemoteV2,
     persist: bool,
     capsys: pytest.CaptureFixture,
     allocation_flags: list[str],
     file_regression: FileRegressionFixture,
     slurm_account_on_cluster: str,
 ):
-    if login_node.hostname == "localhost":
+    if login_node_v2.hostname == "localhost":
         pytest.skip(
             "TODO: This test doesn't yet work with the slurm cluster spun up in the GitHub CI."
         )
 
-    home = await login_node.get_output_async("echo $HOME")
-    scratch = await login_node.get_output_async("echo $SCRATCH")
+    home = await login_node_v2.get_output_async("echo $HOME")
+    scratch = await login_node_v2.get_output_async("echo $SCRATCH")
 
     start = datetime.datetime.now() - timedelta(minutes=5)
     jobs_before = get_recent_jobs_info_dicts(
-        login_node, since=datetime.datetime.now() - start
+        login_node_v2, since=datetime.datetime.now() - start
     )
     jobs_before = {
         int(job_info["JobID"]): job_info
@@ -140,7 +140,7 @@ async def test_code(
             job=None,
             node=None,
             alloc=allocation_flags,
-            cluster=login_node.hostname,  # type: ignore
+            cluster=login_node_v2.hostname,  # type: ignore
         )
 
     # Get the output that was printed while running that command.
@@ -162,12 +162,12 @@ async def test_code(
 
     job_info = await get_job_info(
         job_id=job_id,
-        login_node=login_node,
+        login_node=login_node_v2,
         fields=("JobID", "JobName", "Node", "WorkDir", "State"),
     )
     if node_hostname is None:
         node_hostname = get_hostname_to_use_for_compute_node(
-            job_info["Node"], cluster=login_node.hostname
+            job_info["Node"], cluster=login_node_v2.hostname
         )
     assert node_hostname and node_hostname != "None"