File size: 6,242 Bytes
1b35d41
 
 
 
 
 
364b897
 
 
 
1b35d41
 
 
364b897
1b35d41
364b897
 
1b35d41
 
 
 
364b897
1b35d41
 
 
 
 
 
364b897
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1b35d41
364b897
 
 
 
1b35d41
 
 
 
 
 
364b897
 
 
 
 
 
 
1b35d41
 
 
 
364b897
1b35d41
 
 
 
 
 
 
 
364b897
 
 
 
 
 
 
 
1b35d41
 
364b897
 
 
 
 
 
 
1b35d41
364b897
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1b35d41
364b897
 
 
 
1b35d41
 
 
364b897
 
 
1b35d41
 
364b897
 
 
 
 
1b35d41
 
 
364b897
1b35d41
 
 
 
 
364b897
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# 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)