code_migration/docker_sandbox.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Docker sandbox for running tests inside containers.

Falls back to running tests via venv when Docker is not available.
"""

from __future__ import annotations

import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path


@dataclass
class TestResult:
    """Result of a test execution."""

    truncated_log: str = ""
    full_log: str | None = None
    exit_code: int | None = None


def _is_docker_available() -> bool:
    """Check if Docker CLI exists and daemon is running."""
    if not shutil.which("docker"):
        return False
    try:
        subprocess.run(["docker", "info"], capture_output=True, timeout=5, check=True)
        return True
    except Exception:
        return False


def _truncate_log(full_log: str, max_lines: int = 100) -> str:
    """Truncate log to last N lines with line numbers."""
    lines = full_log.splitlines(keepends=True)
    num_lines = len(lines)
    lines_to_show = lines[-max_lines:]
    start_line = num_lines - len(lines_to_show) + 1

    truncated = "".join(
        f"{i}: {line}"
        for i, line in enumerate(lines_to_show, start=start_line)
    )
    lines_before = max(0, num_lines - max_lines)
    if lines_before > 0:
        truncated = f"({lines_before} lines above)\n" + truncated
    return truncated.rstrip()


class DockerSandbox:
    """Build Docker images and run containers with timeout/memory limits.

    Falls back to running tests via venv when Docker is not available.
    """

    def __init__(self, timeout: int = 600, memory_limit: str = "16g") -> None:
        self.timeout = timeout
        self.memory_limit = memory_limit

    def run_tests(self, image_name: str, workspace_dir: str) -> TestResult:
        """Run tests. Uses Docker if available, otherwise runs via venv."""
        if _is_docker_available():
            return self._run_tests_docker(image_name, workspace_dir)
        else:
            return self._run_tests_venv(image_name, workspace_dir)

    def _run_tests_docker(self, image_name: str, workspace_dir: str) -> TestResult:
        """Run the container, capture output, truncate to last 100 lines."""
        container_name = image_name
        try:
            command = [
                "timeout", "--foreground", "-s", "KILL", str(self.timeout),
                "docker", "run",
                "-v", f"{workspace_dir}:/work",
                f"--memory={self.memory_limit}",
                "--name", container_name,
                image_name,
            ]

            result = subprocess.run(
                command, text=True, stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT, check=False,
            )

            return TestResult(
                truncated_log=_truncate_log(result.stdout),
                full_log=result.stdout,
                exit_code=result.returncode,
            )

        except Exception as e:
            return TestResult(truncated_log=f"Test execution failed: {e}", full_log=None, exit_code=None)
        finally:
            try:
                self.remove_container(container_name)
            except Exception:
                pass

    def _run_tests_venv(self, image_name: str, workspace_dir: str) -> TestResult:
        """Run tests using the venv created during setup (no Docker)."""
        try:
            venv_python = str(Path(workspace_dir) / ".venv" / "bin" / "python")
            if not Path(venv_python).exists():
                return TestResult(
                    truncated_log="No venv found. Setup may have failed.",
                    full_log=None, exit_code=1,
                )

            # Find the test script: image_name is like "brmc__shortbus_new"
            base_name = image_name.replace("_new", "")
            test_script = Path(workspace_dir) / f"test_{base_name}.sh"

            if not test_script.exists():
                candidates = list(Path(workspace_dir).glob("test_*.sh"))
                test_script = candidates[0] if candidates else None

            if not test_script or not test_script.exists():
                return TestResult(
                    truncated_log="No test script found in workspace.",
                    full_log=None, exit_code=1,
                )

            # Parse the test command from the script
            test_cmd = None
            with open(test_script) as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#") or line.startswith("set "):
                        continue
                    test_cmd = line
                    break

            if not test_cmd:
                return TestResult(
                    truncated_log="Test script is empty.",
                    full_log=None, exit_code=1,
                )

            # Replace python with venv python
            test_cmd = test_cmd.replace("python3 ", f"{venv_python} ").replace("python ", f"{venv_python} ")

            print(f"[VENV-TEST] Running: {test_cmd}", flush=True)

            result = subprocess.run(
                test_cmd, shell=True, cwd=workspace_dir,
                text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                timeout=self.timeout,
            )

            return TestResult(
                truncated_log=_truncate_log(result.stdout),
                full_log=result.stdout,
                exit_code=result.returncode,
            )

        except subprocess.TimeoutExpired:
            return TestResult(
                truncated_log=f"Test execution timed out after {self.timeout}s.",
                full_log=None, exit_code=1,
            )
        except Exception as e:
            return TestResult(
                truncated_log=f"Test execution failed: {e}",
                full_log=None, exit_code=None,
            )

    @staticmethod
    def remove_container(container_name: str) -> None:
        """Force-remove a Docker container."""
        subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)