update

Co-authored-by: Copilot <copilot@github.com>

tests/test_cli.py

@@ -1,3 +1,5 @@
+from unittest.mock import AsyncMock, MagicMock
+
 from typer.testing import CliRunner
 
 from tai.cli import app
@@ -5,6 +7,24 @@ from tai.collectors import CollectedItem, CollectionReport
 from tai.ssh_client import SSHCommandResult
 
 
+def _mock_session(
+    monkeypatch,  # type: ignore[no-untyped-def]
+    *,
+    probe_result: SSHCommandResult | None = None,
+    probe_raises: Exception | None = None,
+) -> MagicMock:
+    """Patch SSHClient.connect to return a mock session."""
+    session = MagicMock()
+    session.__aenter__ = AsyncMock(return_value=session)
+    session.__aexit__ = AsyncMock(return_value=None)
+    if probe_raises:
+        session.probe = AsyncMock(side_effect=probe_raises)
+    else:
+        session.probe = AsyncMock(return_value=probe_result)
+    monkeypatch.setattr("tai.cli.SSHClient.connect", lambda _self, **kw: session)
+    return session
+
+
 def test_run_command_prints_scaffold_summary() -> None:
     runner = CliRunner()
     result = runner.invoke(
@@ -25,33 +45,23 @@ def test_run_command_prints_scaffold_summary() -> None:
     )
 
     assert result.exit_code == 0
-    assert "tai scaffold ready" in result.stdout
+    assert "tai" in result.stdout
     assert "host=web01" in result.stdout
     assert "port=5566" in result.stdout
 
 
 def test_probe_success_prints_remote_output_by_default(monkeypatch) -> None:  # type: ignore[no-untyped-def]
-    async def fake_probe(self) -> SSHCommandResult:  # type: ignore[no-untyped-def]
-        return SSHCommandResult(
-            command="uname -a",
-            exit_code=0,
-            stdout="Linux ssh 6.12.0",
-            stderr="",
-        )
-
-    monkeypatch.setattr("tai.cli.SSHClient.probe", fake_probe)
+    _mock_session(
+        monkeypatch,
+        probe_result=SSHCommandResult(
+            command="uname -a", exit_code=0, stdout="Linux ssh 6.12.0", stderr=""
+        ),
+    )
 
     runner = CliRunner()
     result = runner.invoke(
         app,
-        [
-            "apache failed",
-            "--host",
-            "ssh.archflux.net",
-            "--port",
-            "5566",
-            "--probe",
-        ],
+        ["apache failed", "--host", "ssh.archflux.net", "--port", "5566", "--probe"],
     )
 
     assert result.exit_code == 0
@@ -60,27 +70,20 @@ def test_probe_success_prints_remote_output_by_default(monkeypatch) -> None:  #
 
 
 def test_probe_failure_returns_non_zero(monkeypatch) -> None:  # type: ignore[no-untyped-def]
-    async def fake_probe(self) -> SSHCommandResult:  # type: ignore[no-untyped-def]
-        return SSHCommandResult(
+    _mock_session(
+        monkeypatch,
+        probe_result=SSHCommandResult(
             command="uname -a",
             exit_code=255,
             stdout="",
             stderr="Permission denied (publickey,password).",
-        )
-
-    monkeypatch.setattr("tai.cli.SSHClient.probe", fake_probe)
+        ),
+    )
 
     runner = CliRunner()
     result = runner.invoke(
         app,
-        [
-            "apache failed",
-            "--host",
-            "ssh.archflux.net",
-            "--port",
-            "5566",
-            "--probe",
-        ],
+        ["apache failed", "--host", "ssh.archflux.net", "--port", "5566", "--probe"],
     )
 
     assert result.exit_code == 1
@@ -88,7 +91,9 @@ def test_probe_failure_returns_non_zero(monkeypatch) -> None:  # type: ignore[no
 
 
 def test_collect_success_prints_summary(monkeypatch) -> None:  # type: ignore[no-untyped-def]
-    async def fake_collect_from_plan(_client, _plan) -> CollectionReport:  # type: ignore[no-untyped-def]
+    _mock_session(monkeypatch)
+
+    async def fake_collect_from_plan(_session, _plan) -> CollectionReport:  # type: ignore[no-untyped-def]
         return CollectionReport(
             host="ssh.archflux.net",
             items=[