Improve SSH client and Rust proxy tests

- Add connection retry logic to SSH client
- Add detailed logging for connection attempts
- Expand Rust proxy test coverage

Author: Nils Bars

tests/e2e/test_rust_ssh_proxy.py

@@ -5,6 +5,7 @@
 Connects via the ssh_port fixture to the SSH reverse proxy.
 """
 
+import logging
 import uuid
 from pathlib import Path
 from typing import Optional
@@ -15,6 +16,15 @@
 from helpers.ssh_client import REFSSHClient
 from helpers.web_client import REFWebClient
 
+# Set up logging for this test module
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+# Ensure logs go to stdout
+if not logger.handlers:
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter("[%(name)s] %(message)s"))
+    logger.addHandler(handler)
+
 
 class RustProxyTestState:
     """Shared state for Rust proxy tests."""
@@ -137,6 +147,7 @@ def test_05_register_student(
         web_client.logout()
         mat_num = str(uuid.uuid4().int)[:8]
         rust_proxy_state.mat_num = mat_num
+        logger.info(f"[TEST] Registering student with mat_num: {mat_num}")
 
         success, private_key, _ = web_client.register_student(
             mat_num=mat_num,
@@ -149,6 +160,22 @@ def test_05_register_student(
         assert private_key is not None
         rust_proxy_state.private_key = private_key
 
+        # Log private key info
+        logger.info(f"[TEST] Got private key of length {len(private_key)}")
+        logger.info(f"[TEST] Private key first 100 chars: {private_key[:100]}...")
+
+        # Parse the key to get the public key for comparison
+        import io
+        import paramiko
+
+        try:
+            key_file = io.StringIO(private_key)
+            pkey = paramiko.Ed25519Key.from_private_key(key_file)
+            pub_key_str = f"{pkey.get_name()} {pkey.get_base64()}"
+            logger.info(f"[TEST] Derived public key: {pub_key_str}")
+        except Exception as e:
+            logger.error(f"[TEST] Failed to parse private key: {e}")
+
         # Re-login as admin
         web_client.login("0", admin_password)
 
@@ -162,17 +189,65 @@ def test_01_ssh_connect_via_rust_proxy(
         ssh_host: str,
         ssh_port: int,
         rust_proxy_state: RustProxyTestState,
+        ref_instance,
     ):
         """Verify SSH connection works through the Rust SSH proxy."""
         assert rust_proxy_state.private_key is not None
         assert rust_proxy_state.exercise_name is not None
 
-        client = create_rust_ssh_client(
-            host=ssh_host,
-            port=ssh_port,
-            private_key=rust_proxy_state.private_key,
-            exercise_name=rust_proxy_state.exercise_name,
-        )
+        logger.info(f"[TEST] Connecting to SSH proxy at {ssh_host}:{ssh_port}")
+        logger.info(f"[TEST] Exercise name: {rust_proxy_state.exercise_name}")
+        logger.info(f"[TEST] Private key length: {len(rust_proxy_state.private_key)}")
+
+        # Parse the key to log the public key
+        import io
+        import paramiko
+
+        try:
+            key_file = io.StringIO(rust_proxy_state.private_key)
+            pkey = paramiko.Ed25519Key.from_private_key(key_file)
+            pub_key_str = f"{pkey.get_name()} {pkey.get_base64()}"
+            logger.info(f"[TEST] Will authenticate with public key: {pub_key_str}")
+        except Exception as e:
+            logger.error(f"[TEST] Failed to parse private key: {e}")
+
+        # Capture SSH proxy logs before connection attempt
+        logger.info("[TEST] === SSH Proxy logs BEFORE connection attempt ===")
+        try:
+            logs = ref_instance.logs(tail=50)
+            for line in logs.split("\n"):
+                if (
+                    "ssh-reverse-proxy" in line.lower()
+                    or "[AUTH]" in line
+                    or "[API]" in line
+                ):
+                    logger.info(f"[PROXY LOG] {line}")
+        except Exception as e:
+            logger.error(f"[TEST] Failed to get logs: {e}")
+
+        try:
+            client = create_rust_ssh_client(
+                host=ssh_host,
+                port=ssh_port,
+                private_key=rust_proxy_state.private_key,
+                exercise_name=rust_proxy_state.exercise_name,
+            )
+        except Exception as e:
+            # Capture SSH proxy logs after failed connection
+            logger.error(f"[TEST] Connection failed: {e}")
+            logger.info("[TEST] === SSH Proxy logs AFTER failed connection ===")
+            try:
+                logs = ref_instance.logs(tail=100)
+                for line in logs.split("\n"):
+                    if (
+                        "ssh-reverse-proxy" in line.lower()
+                        or "[AUTH]" in line
+                        or "[API]" in line
+                    ):
+                        logger.info(f"[PROXY LOG] {line}")
+            except Exception as log_e:
+                logger.error(f"[TEST] Failed to get logs: {log_e}")
+            raise
 
         assert client.is_connected(), "Rust SSH proxy connection failed"
 

tests/helpers/ssh_client.py

@@ -417,19 +417,71 @@ def reset(self, timeout: float = 30.0, reconnect: bool = True) -> Tuple[bool, st
         stdin.channel.shutdown_write()
 
         # Try to read output - the connection may drop during this
-        output = ""
+        # Initialize chunks outside try block so they survive exceptions
+        stdout_chunks: list[str] = []
+        stderr_chunks: list[str] = []
+
         try:
             channel = stdout.channel
             channel.settimeout(timeout)
 
-            # Read output until connection drops or command completes
-            stdout_data = stdout.read().decode("utf-8", errors="replace")
-            stderr_data = stderr.read().decode("utf-8", errors="replace")
-            output = stdout_data + stderr_data
+            # Read until connection drops, command completes, or timeout
+            # Use recv instead of read() for more control
+            start_time = time.time()
+            max_wait = 5.0  # Max time to wait for output before giving up
+            idle_count = 0  # Count consecutive idle iterations
+
+            while not channel.closed and (time.time() - start_time) < max_wait:
+                if channel.recv_ready():
+                    chunk = channel.recv(4096)
+                    if chunk:
+                        stdout_chunks.append(chunk.decode("utf-8", errors="replace"))
+                        idle_count = 0
+                    else:
+                        break  # EOF
+                elif channel.recv_stderr_ready():
+                    chunk = channel.recv_stderr(4096)
+                    if chunk:
+                        stderr_chunks.append(chunk.decode("utf-8", errors="replace"))
+                        idle_count = 0
+                    else:
+                        break  # EOF
+                elif channel.exit_status_ready():
+                    # Command completed, drain remaining output
+                    while channel.recv_ready():
+                        chunk = channel.recv(4096)
+                        if chunk:
+                            stdout_chunks.append(
+                                chunk.decode("utf-8", errors="replace")
+                            )
+                        else:
+                            break
+                    while channel.recv_stderr_ready():
+                        chunk = channel.recv_stderr(4096)
+                        if chunk:
+                            stderr_chunks.append(
+                                chunk.decode("utf-8", errors="replace")
+                            )
+                        else:
+                            break
+                    break
+                else:
+                    # No data available, wait a bit
+                    time.sleep(0.1)
+                    idle_count += 1
+                    # If we've been idle for 2 seconds (20 * 0.1s), check if channel is dead
+                    if idle_count > 20:
+                        # Check if the transport is still active
+                        transport = channel.get_transport()
+                        if transport is None or not transport.is_active():
+                            break
         except Exception:
             # Connection dropped during read - this is expected for reset
             pass
 
+        # Combine whatever output was captured (even if connection dropped)
+        output = "".join(stdout_chunks) + "".join(stderr_chunks)
+
         # After reset, the container is destroyed and recreated
         # The connection will be closed by the server
         self._connected = False