Add tests for the remote debug via SSH command

tests/commands/execution/test_ssh.py

@@ -0,0 +1,68 @@
+import pathlib
+import subprocess
+import time
+
+from tests.commands.execution.test_watch import no_sleep
+from tests.commands.execution.utils import get_execution_data_mock
+from tests.fixture_data import EXECUTION_DATA
+from valohai_cli.commands.execution.ssh import ssh
+
+
+def test_ssh_in_completed_execution(runner, logged_in_and_linked, monkeypatch):
+    with get_execution_data_mock():
+        counter = EXECUTION_DATA["counter"]
+        output = runner.invoke(ssh, [str(counter)], catch_exceptions=False).output
+        assert f"Error: Execution #{counter} is complete. Cannot SSH into it.\n" in output
+
+        EXECUTION_DATA["status"] = "queued"
+        monkeypatch.setattr(time, "sleep", no_sleep)
+
+
+def test_ssh_in_queued_execution(runner, logged_in_and_linked, monkeypatch):
+    counter = EXECUTION_DATA["counter"]
+    EXECUTION_DATA["status"] = "queued"
+    monkeypatch.setattr(time, "sleep", no_sleep)
+    with get_execution_data_mock():
+        output = runner.invoke(ssh, [str(counter)], catch_exceptions=False).output
+        assert f"Execution #{counter} is queued. Waiting for it to start...\n" in output
+
+
+def test_ssh_with_no_ssh_details_present(runner, logged_in_and_linked, monkeypatch):
+    counter = EXECUTION_DATA["counter"]
+    EXECUTION_DATA["status"] = "started"
+    monkeypatch.setattr(time, "sleep", lambda x: None)
+    monkeypatch.setattr("valohai_cli.utils.ssh.find_ssh_details", lambda x: None)
+    with get_execution_data_mock():
+        output = runner.invoke(ssh, [str(counter)], catch_exceptions=False).output
+        assert "1/5 Retrying: No SSH details found...\n" in output
+        assert "2/5 Retrying: No SSH details found...\n" in output
+        assert "3/5 Retrying: No SSH details found...\n" in output
+        assert "4/5 Retrying: No SSH details found...\n" in output
+        assert "5/5 Retrying: No SSH details found...\n" in output
+
+
+def test_ssh(runner, logged_in_and_linked, monkeypatch):
+    counter = EXECUTION_DATA["counter"]
+    EXECUTION_DATA["status"] = "started"
+    monkeypatch.setattr(time, "sleep", lambda x: None)
+
+    def mock_prompt():
+        return pathlib.Path("/path/to/private")
+
+    monkeypatch.setattr(
+        "valohai_cli.commands.execution.ssh.select_private_key_from_possible_directories",
+        mock_prompt,
+    )
+
+    with get_execution_data_mock():
+        output = runner.invoke(ssh, [str(counter)], catch_exceptions=False).output
+        assert "SSH address is 127.0.0.1:2222" in output
+
+        def mock_ssh_connection(*args, **kwargs):
+            print("Mock SSH interactive session executed")
+            return subprocess.CompletedProcess(args=args, returncode=0, stderr=b"")
+
+        monkeypatch.setattr("valohai_cli.commands.execution.ssh.make_ssh_connection", mock_ssh_connection)
+
+        output = runner.invoke(ssh, [str(counter)], input="1", catch_exceptions=False).output
+        assert "Mock SSH interactive session executed" in output

tests/commands/execution/utils.py

@@ -10,6 +10,7 @@
     OUTPUT_DATUM_DOWNLOAD_RESPONSE_DATA,
     OUTPUT_DATUM_RESPONSE_DATA,
     PROJECT_DATA,
+    STATUS_EVENT_RESPONSE_DATA,
 )
 
 API_PREFIX = "https://app.valohai.com/api/v0/"
@@ -29,6 +30,7 @@ def get_execution_data_mock():
     m.get(f"{API_PREFIX}executions/", json={"results": [EXECUTION_DATA]})
     m.get(f"{API_PREFIX}executions/{exec_id}/", json=EXECUTION_DATA)
     m.get(f"{API_PREFIX}executions/{exec_id}/events/", json=EVENT_RESPONSE_DATA)
+    m.get(f"{API_PREFIX}executions/{exec_id}/status-events/?limit=100", json=STATUS_EVENT_RESPONSE_DATA)
     m.get(f"{API_PREFIX}data/?output_execution={exec_id}&limit=9000", json=OUTPUT_DATUM_RESPONSE_DATA)
     m.get(f"{API_PREFIX}data/{datum_id}/download/", json=OUTPUT_DATUM_DOWNLOAD_RESPONSE_DATA)
     execution_by_counter_url = f"{API_PREFIX}executions/{project_id}:{execution_counter}/"

tests/fixture_data.py

@@ -330,7 +330,6 @@
         path: model.h5
 """
 
-
 PIPELINE_YAML = (
     CONFIG_YAML
     + """
@@ -665,3 +664,19 @@ def main(config) -> Pipeline:
         description: Model output file from TensorFlow
         path: model.h5
 """
+
+STATUS_EVENT_RESPONSE_DATA = {
+    "total": 2,
+    "status_events": [
+        {
+            "stream": "status",
+            "message": '::ssh::{"port": 2222, "address": "127.0.0.1"}',
+            "time": "2024-09-04T12:16:20.722000",
+        },
+        {
+            "stream": "status",
+            "message": "   $ ssh -i <path-to-private-key> 127.0.0.1 -p 2222 -t /bin/bash",
+            "time": "2024-09-04T12:16:20.723000",
+        },
+    ],
+}