mcpbr 0.4.13__tar.gz → 0.4.14__tar.gz

{mcpbr-0.4.13 → mcpbr-0.4.14}/.claude-plugin/marketplace.json

@@ -1,7 +1,7 @@
 {
   "$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
   "name": "mcpbr",
-  "version": "0.4.13",
+  "version": "0.4.14",
   "description": "mcpbr - MCP Benchmark Runner plugin marketplace",
   "owner": {
     "name": "mcpbr Contributors",
@@ -11,7 +11,7 @@
     {
       "name": "mcpbr",
       "description": "Expert benchmark runner for MCP servers using mcpbr. Handles Docker checks, config generation, and result parsing.",
-      "version": "0.4.13",
+      "version": "0.4.14",
       "author": {
         "name": "mcpbr Contributors"
       },

{mcpbr-0.4.13 → mcpbr-0.4.14}/.claude-plugin/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@greynewell/mcpbr-claude-plugin",
-  "version": "0.4.13",
+  "version": "0.4.14",
   "description": "Claude Code plugin for mcpbr - Expert benchmark runner for MCP servers with specialized skills",
   "keywords": [
     "claude-code",

{mcpbr-0.4.13 → mcpbr-0.4.14}/.claude-plugin/plugin.json

@@ -1,6 +1,6 @@
 {
   "name": "mcpbr",
-  "version": "0.4.13",
+  "version": "0.4.14",
   "description": "Expert benchmark runner for MCP servers using mcpbr. Handles Docker checks, config generation, and result parsing.",
   "schema_version": "1.0"
 }

{mcpbr-0.4.13 → mcpbr-0.4.14}/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- **Repository-aware test commands for non-pytest projects** (#365): Use upstream SWE-bench test command specs for sympy (`bin/test`), sphinx (`tox`), and other non-pytest repos instead of defaulting to `python -m pytest`
+- **Flaky Azure and trial mode tests**: Fixed tests that depended on local `~/.ssh/mcpbr_azure` state and updated assertions for multi-step dependency installation
 - **SEO improvements** for documentation site
   - Added robots.txt with sitemap reference
   - Added Open Graph and Twitter Card meta tags on all pages

{mcpbr-0.4.13 → mcpbr-0.4.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mcpbr
-Version: 0.4.13
+Version: 0.4.14
 Summary: Model Context Protocol Benchmark Runner - evaluate MCP servers against software engineering benchmarks
 Project-URL: Homepage, https://github.com/greynewell/mcpbr
 Project-URL: Repository, https://github.com/greynewell/mcpbr

{mcpbr-0.4.13 → mcpbr-0.4.14}/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@greynewell/mcpbr",
-  "version": "0.4.13",
+  "version": "0.4.14",
   "description": "Model Context Protocol Benchmark Runner - CLI tool for evaluating MCP servers",
   "keywords": [
     "mcpbr",

{mcpbr-0.4.13 → mcpbr-0.4.14}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mcpbr"
-version = "0.4.13"
+version = "0.4.14"
 description = "Model Context Protocol Benchmark Runner - evaluate MCP servers against software engineering benchmarks"
 readme = "README.md"
 license = "MIT"

{mcpbr-0.4.13 → mcpbr-0.4.14}/src/mcpbr/config.py

@@ -109,6 +109,16 @@ class MCPServerConfig(BaseModel):
         default=900000,
         description="Timeout in milliseconds for MCP tool execution (default: 15 min for long-running tools)",
     )
+    setup_command: str | None = Field(
+        default=None,
+        description="Shell command to run inside the container BEFORE the agent starts. "
+        "Runs outside the task timer (does not count against timeout_seconds). "
+        "Use {workdir} as placeholder. Useful for pre-computing caches.",
+    )
+    setup_timeout_ms: int = Field(
+        default=900000,
+        description="Timeout in milliseconds for the setup_command (default: 15 min)",
+    )
 
     def get_args_for_workdir(self, workdir: str) -> list[str]:
         """Replace {workdir} placeholder in args with actual path."""
@@ -117,6 +127,12 @@ class MCPServerConfig(BaseModel):
             result.append(arg.replace("{workdir}", workdir))
         return result
 
+    def get_setup_command_for_workdir(self, workdir: str) -> str | None:
+        """Replace {workdir} placeholder in setup_command with actual path."""
+        if self.setup_command is None:
+            return None
+        return self.setup_command.replace("{workdir}", workdir)
+
     def get_expanded_env(self) -> dict[str, str]:
         """Expand ${VAR} references in env values using os.environ.
 
@@ -400,6 +416,12 @@ class HarnessConfig(BaseModel):
         description="Enable comprehensive performance profiling (tool latency, memory, overhead)",
     )
 
+    volumes: dict[str, str] = Field(
+        default_factory=dict,
+        description="Additional volume mounts (read-write) for Docker containers (host_path: container_path). "
+        "Mounted into every container, persists across tasks. Useful for pre-computed caches.",
+    )
+
     infrastructure: InfrastructureConfig = Field(
         default_factory=InfrastructureConfig,
         description="Infrastructure configuration (local or azure)",

{mcpbr-0.4.13 → mcpbr-0.4.14}/src/mcpbr/docker_env.py

@@ -314,14 +314,18 @@ class DockerEnvironmentManager:
     FALLBACK_IMAGE = "mcpbr-env"
     DOCKERFILE_PATH = Path(__file__).parent.parent.parent / "Dockerfile"
 
-    def __init__(self, use_prebuilt: bool = True) -> None:
+    def __init__(
+        self, use_prebuilt: bool = True, extra_volumes: dict[str, str] | None = None
+    ) -> None:
         """Initialize the Docker environment manager.
 
         Args:
             use_prebuilt: If True, try to use pre-built SWE-bench images first.
+            extra_volumes: Additional volume mounts (read-write) (host_path -> container_path).
         """
         self.client = docker.from_env()
         self.use_prebuilt = use_prebuilt
+        self._extra_volumes = extra_volumes or {}
         self._fallback_image_built = False
         self._temp_dirs: list[tempfile.TemporaryDirectory[str]] = []
         self._containers: list[Container] = []
@@ -488,6 +492,15 @@ CMD ["/bin/bash"]
 
             for attempt in range(max_retries + 1):
                 try:
+                    volumes_dict: dict[str, dict[str, str]] = {
+                        host_workdir: {"bind": "/workspace", "mode": "rw"},
+                    }
+                    for host_path, container_path in self._extra_volumes.items():
+                        volumes_dict[os.path.abspath(host_path)] = {
+                            "bind": container_path,
+                            "mode": "rw",
+                        }
+
                     container = self.client.containers.run(
                         image_name,
                         command="tail -f /dev/null",
@@ -495,9 +508,7 @@ CMD ["/bin/bash"]
                         detach=True,
                         platform="linux/amd64" if uses_prebuilt else None,
                         network_mode="bridge",  # Enable network for API calls
-                        volumes={
-                            host_workdir: {"bind": "/workspace", "mode": "rw"},
-                        },
+                        volumes=volumes_dict,
                         working_dir=container_workdir,
                         remove=False,
                         labels={

{mcpbr-0.4.13 → mcpbr-0.4.14}/src/mcpbr/evaluation.py

@@ -137,6 +137,7 @@ async def run_tests(
     timeout: int = 120,
     uses_prebuilt: bool = False,
     workdir: str | None = None,
+    repo: str | None = None,
 ) -> TestResults:
     """Run a list of tests and return results.
 
@@ -146,6 +147,7 @@ async def run_tests(
         timeout: Timeout per test in seconds.
         uses_prebuilt: Whether a pre-built SWE-bench image is being used.
         workdir: Working directory to run tests from. Defaults to env.workdir.
+        repo: Repository identifier for looking up the correct test runner.
 
     Returns:
         TestResults with pass/fail counts.
@@ -157,7 +159,7 @@ async def run_tests(
     passed = 0
 
     for test in tests:
-        test_cmd = _build_test_command(test, uses_prebuilt)
+        test_cmd = _build_test_command(test, uses_prebuilt, repo=repo)
 
         try:
             exit_code, stdout, stderr = await env.exec_command(
@@ -198,7 +200,7 @@ async def run_tests(
     )
 
 
-def _build_test_command(test: str, uses_prebuilt: bool = False) -> str:
+def _build_test_command(test: str, uses_prebuilt: bool = False, repo: str | None = None) -> str:
     """Build a test command for the given test identifier.
 
     Args:
@@ -206,18 +208,29 @@ def _build_test_command(test: str, uses_prebuilt: bool = False) -> str:
             - pytest: "tests/test_file.py::test_func" or "tests/test_file.py"
             - Django: "test_method (module.TestClass)" or "module.tests.TestClass.test_method"
         uses_prebuilt: If True, activate the testbed conda environment first.
+        repo: Repository identifier (e.g., "sympy/sympy") for looking up
+            the correct test runner from upstream SWE-bench specs.
 
     Returns:
         Shell command string to run the test.
     """
     import re
 
+    from .swebench_test_specs import get_repo_test_command
+
     # Pre-built SWE-bench images use a conda environment called 'testbed'
     if uses_prebuilt:
         activate = "source /opt/miniconda3/etc/profile.d/conda.sh && conda activate testbed && "
     else:
         activate = ""
 
+    # Check upstream SWE-bench test command mapping for non-pytest runners
+    if repo:
+        upstream_cmd = get_repo_test_command(repo)
+        if upstream_cmd and "runtests.py" not in upstream_cmd and "pytest" not in upstream_cmd:
+            # Non-pytest, non-Django project (e.g., sympy uses bin/test)
+            return f"{activate}{upstream_cmd} {test}"
+
     # Detect Django test format: "test_method (module.TestClass)"
     if "(" in test and ")" in test and "." in test:
         # Extract module path from parentheses
@@ -344,12 +357,15 @@ async def evaluate_patch(
     if not env.uses_prebuilt:
         await _install_dependencies(env)
 
+    repo = task.get("repo")
+
     fail_to_pass_results = await run_tests(
         env,
         fail_to_pass_tests,
         timeout=test_timeout,
         uses_prebuilt=env.uses_prebuilt,
         workdir=eval_workdir,
+        repo=repo,
     )
 
     pass_to_pass_results = await run_tests(
@@ -358,6 +374,7 @@ async def evaluate_patch(
         timeout=test_timeout,
         uses_prebuilt=env.uses_prebuilt,
         workdir=eval_workdir,
+        repo=repo,
     )
 
     resolved = (

{mcpbr-0.4.13 → mcpbr-0.4.14}/src/mcpbr/harness.py

@@ -962,7 +962,10 @@ async def run_evaluation(
                 "args": config.mcp_server.args if config.mcp_server else [],
             }
 
-    docker_manager = DockerEnvironmentManager(use_prebuilt=config.use_prebuilt_images)
+    docker_manager = DockerEnvironmentManager(
+        use_prebuilt=config.use_prebuilt_images,
+        extra_volumes=config.volumes,
+    )
 
     results: list[TaskResult] = []
     # Add cached results if using state tracker

{mcpbr-0.4.13 → mcpbr-0.4.14}/src/mcpbr/harnesses.py

@@ -895,6 +895,35 @@ class ClaudeCodeHarness:
                     cost_usd=None,
                 )
 
+        # Run setup_command if configured (BEFORE agent, OUTSIDE task timer).
+        # This is the right place for expensive one-time operations like
+        # pre-computing caches that should not count against timeout_seconds.
+        if self.mcp_server and self.mcp_server.setup_command:
+            setup_cmd = self.mcp_server.get_setup_command_for_workdir(env.workdir)
+            setup_timeout = int(self.mcp_server.setup_timeout_ms / 1000)
+
+            if verbose:
+                self._console.print(
+                    f"[cyan]Running setup command (timeout: {setup_timeout:.0f}s)...[/cyan]"
+                )
+
+            setup_full_cmd = f"source {shlex.quote(env_file)} && {setup_cmd}"
+            setup_exit, _setup_stdout, setup_stderr = await env.exec_command(
+                ["/bin/bash", "-c", setup_full_cmd],
+                timeout=setup_timeout,
+            )
+
+            if setup_exit != 0:
+                if verbose:
+                    self._console.print(
+                        f"[yellow]⚠ Setup command exited with code {setup_exit}[/yellow]"
+                    )
+                    if setup_stderr:
+                        self._console.print(f"[dim]{setup_stderr[:500]}[/dim]")
+                # Non-fatal: continue with agent even if setup fails
+            elif verbose:
+                self._console.print("[green]✓ Setup command completed[/green]")
+
         try:
             claude_args = [
                 "--print",

mcpbr-0.4.14/src/mcpbr/swebench_test_specs.py

@@ -0,0 +1,33 @@
+"""Test command specs from upstream SWE-bench harness.
+
+Maps repositories to their correct test commands. mcpbr defaults to pytest
+for all non-Django projects, but some projects (e.g., sympy) use custom test
+runners that aren't pytest-compatible.
+
+Source: https://github.com/SWE-bench/SWE-bench/blob/main/swebench/harness/constants/python.py
+"""
+
+# Base test commands per framework (from upstream constants/python.py)
+TEST_PYTEST = "pytest --no-header -rA --tb=no -p no:cacheprovider"
+TEST_DJANGO = "./tests/runtests.py --verbosity 2 --settings=test_sqlite --parallel 1"
+TEST_SYMPY = "PYTHONWARNINGS='ignore::UserWarning,ignore::SyntaxWarning' bin/test -C --verbose"
+TEST_SPHINX = "tox --current-env -epy39 -v --"
+TEST_ASTROPY = "pytest -rA -vv -o console_output_style=classic --tb=no"
+TEST_SEABORN = "pytest --no-header -rA"
+
+# Repo → test command mapping
+# Only non-pytest entries need to be here — pytest is the default fallback.
+# Django is included for documentation but its existing handler takes precedence.
+REPO_TO_TEST_CMD: dict[str, str] = {
+    "sympy/sympy": TEST_SYMPY,
+    "django/django": TEST_DJANGO,
+    "sphinx-doc/sphinx": TEST_SPHINX,
+}
+
+
+def get_repo_test_command(repo: str) -> str | None:
+    """Look up the upstream test command for a repo.
+
+    Returns None if repo uses standard pytest (handled by existing logic).
+    """
+    return REPO_TO_TEST_CMD.get(repo)

{mcpbr-0.4.13 → mcpbr-0.4.14}/tests/infrastructure/test_azure.py

@@ -155,11 +155,15 @@ class TestVMProvisioning:
         mock_time: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test successful VM creation."""
-        # Mock ssh-keygen, resource group show (exists), vm create
+        # Use existing SSH key to avoid depending on ~/.ssh/mcpbr_azure state
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
+
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show (exists)
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
         ]
@@ -179,11 +183,15 @@ class TestVMProvisioning:
         mock_time: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test VM creation with resource group creation."""
-        # Mock resource group doesn't exist, then create it
+        # Use existing SSH key to avoid depending on ~/.ssh/mcpbr_azure state
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
+
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=1, stderr="ResourceGroupNotFound"),  # az group show (not found)
             Mock(returncode=0),  # az group create
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
@@ -198,16 +206,19 @@ class TestVMProvisioning:
         self,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test VM creation with SSH key generation."""
-        # Mock ssh-keygen, resource group show, and vm creation
-        mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
-            Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
-            Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
-        ]
+        # Redirect Path.home to tmp_path so ~/.ssh/mcpbr_azure doesn't exist
+        with patch("mcpbr.infrastructure.azure.Path.home", return_value=tmp_path):
+            # Mock ssh-keygen, resource group show, and vm creation
+            mock_run.side_effect = [
+                Mock(returncode=0),  # ssh-keygen
+                Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
+                Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
+            ]
 
-        await azure_provider._create_vm("Standard_D8s_v3")
+            await azure_provider._create_vm("Standard_D8s_v3")
 
         # Verify ssh-keygen was called
         ssh_keygen_call = mock_run.call_args_list[0]
@@ -218,11 +229,15 @@ class TestVMProvisioning:
         self,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test VM creation failure (quota exceeded)."""
-        # Mock ssh-keygen success, resource group show, VM creation failure
+        # Use existing SSH key to avoid depending on ~/.ssh/mcpbr_azure state
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
+
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=1, stderr="QuotaExceeded: Core quota exceeded"),  # az vm create
         ]
@@ -577,13 +592,18 @@ class TestSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test full setup flow (create VM, wait SSH, get IP, install, config, test)."""
         mock_env_get.return_value = "test-api-key"
 
-        # Mock ssh-keygen, resource group show, vm create, vm show
+        # Use existing SSH key to avoid depending on ~/.ssh/mcpbr_azure state
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
+
+        # Mock resource group show, vm create, vm show (no ssh-keygen needed)
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show (note: quoted string in JSON)
@@ -618,11 +638,16 @@ class TestSetup:
         mock_time: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test setup failure rolls back VM creation."""
-        # Mock ssh-keygen success, resource group show, VM creation success, IP retrieval failure
+        # Use existing SSH key to avoid depending on ~/.ssh/mcpbr_azure state
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
+
+        # Mock resource group show, VM creation success, IP retrieval failure
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=1, stderr="VM not found"),  # az vm show (failure)
@@ -687,11 +712,12 @@ class TestSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test setup with generated SSH key."""
         mock_env_get.return_value = "test-api-key"
 
-        # No SSH key configured
+        # No SSH key configured - redirect home to tmp_path so key doesn't exist
         azure_provider.azure_config.ssh_key_path = None
 
         mock_run.side_effect = [
@@ -717,7 +743,8 @@ class TestSetup:
         mock_sftp = MagicMock()
         mock_client.open_sftp.return_value = mock_sftp
 
-        await azure_provider.setup()
+        with patch("mcpbr.infrastructure.azure.Path.home", return_value=tmp_path):
+            await azure_provider.setup()
 
         # Verify ssh-keygen was called
         ssh_keygen_call = mock_run.call_args_list[0]
@@ -793,12 +820,13 @@ class TestEnvironmentSetup:
 
         await azure_provider._install_dependencies()
 
-        # Verify command was executed
-        mock_client.exec_command.assert_called_once()
-        cmd = mock_client.exec_command.call_args[0][0]
-        assert "apt-get update" in cmd
-        assert "docker" in cmd.lower()
-        assert "pip3 install mcpbr" in cmd
+        # Verify all 4 steps were executed (Docker, Python, Node.js, mcpbr)
+        assert mock_client.exec_command.call_count == 4
+        all_cmds = [call[0][0] for call in mock_client.exec_command.call_args_list]
+        all_cmds_str = " ".join(all_cmds)
+        assert "apt-get update" in all_cmds_str
+        assert "docker" in all_cmds_str.lower()
+        assert "pip install mcpbr" in all_cmds_str
 
     async def test_install_dependencies_handles_failures_gracefully(
         self,
@@ -819,7 +847,8 @@ class TestEnvironmentSetup:
         # Should not raise - just log warning
         await azure_provider._install_dependencies()
 
-        mock_client.exec_command.assert_called_once()
+        # All 4 steps still execute even if individual steps fail
+        assert mock_client.exec_command.call_count == 4
 
     async def test_install_dependencies_installs_docker(
         self,
@@ -839,8 +868,9 @@ class TestEnvironmentSetup:
 
         await azure_provider._install_dependencies()
 
-        cmd = mock_client.exec_command.call_args[0][0]
-        assert "get.docker.com" in cmd
+        # Docker install is the first step
+        all_cmds = [call[0][0] for call in mock_client.exec_command.call_args_list]
+        assert any("get.docker.com" in cmd for cmd in all_cmds)
 
     async def test_install_dependencies_installs_python_version(
         self,
@@ -879,8 +909,9 @@ class TestEnvironmentSetup:
 
         await azure_provider._install_dependencies()
 
-        cmd = mock_client.exec_command.call_args[0][0]
-        assert "pip3 install mcpbr" in cmd
+        # mcpbr install is the last step
+        all_cmds = [call[0][0] for call in mock_client.exec_command.call_args_list]
+        assert any("pip install mcpbr" in cmd for cmd in all_cmds)
 
 
 # ============================================================================
@@ -1209,13 +1240,16 @@ class TestUpdatedSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test full setup flow includes dependency installation."""
         mock_env_get.return_value = "test-api-key"
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
 
-        # Mock subprocess calls
+        # Mock subprocess calls (no ssh-keygen needed with existing key)
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show
@@ -1259,12 +1293,15 @@ class TestUpdatedSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test full setup flow includes config transfer."""
         mock_env_get.return_value = "test-api-key"
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
 
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show
@@ -1301,12 +1338,15 @@ class TestUpdatedSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test full setup flow includes env var export."""
         mock_env_get.return_value = "test-api-key"
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
 
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show
@@ -1343,12 +1383,15 @@ class TestUpdatedSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test full setup flow includes test task."""
         mock_env_get.return_value = "test-api-key"
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
 
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show
@@ -1385,12 +1428,15 @@ class TestUpdatedSetup:
         mock_ssh_client: MagicMock,
         mock_run: MagicMock,
         azure_provider: AzureProvider,
+        tmp_path: Path,
     ) -> None:
         """Test setup fails if test task fails."""
         mock_env_get.return_value = "test-api-key"
+        ssh_key = tmp_path / "test_key"
+        ssh_key.touch()
+        azure_provider.azure_config.ssh_key_path = ssh_key
 
         mock_run.side_effect = [
-            Mock(returncode=0),  # ssh-keygen
             Mock(returncode=0, stdout='{"id": "rg-id"}'),  # az group show
             Mock(returncode=0, stdout='{"id": "vm-id"}'),  # az vm create
             Mock(returncode=0, stdout='"1.2.3.4"'),  # az vm show