test_server.py

import subprocess
import unittest
from unittest.mock import Mock, call, patch

from server import MAX_OUTPUT_BYTES, exec_vps


def fake_process(
    stdout: bytes = b"", stderr: bytes = b"", returncode: int = 0
) -> tuple[Mock, object]:
    process = Mock(returncode=returncode)
    process.wait.return_value = returncode
    process.poll.return_value = returncode

    def start(_command: str, **kwargs: object) -> Mock:
        kwargs["stdout"].write(stdout)  # type: ignore[union-attr]
        kwargs["stderr"].write(stderr)  # type: ignore[union-attr]
        return process

    return process, start


class ExecVpsTests(unittest.TestCase):
    def test_returns_structured_stdout_stderr_and_exit_code(self) -> None:
        process, start = fake_process(b"output\n", b"warning\n", 7)

        with patch("server.subprocess.Popen", side_effect=start) as popen:
            result = exec_vps("example", 30)

        self.assertEqual(
            result,
            {
                "exit_code": 7,
                "stdout": "output\n",
                "stderr": "warning\n",
                "timed_out": False,
                "truncated": False,
            },
        )
        self.assertEqual(popen.call_args.args, ("example",))
        self.assertTrue(popen.call_args.kwargs["shell"])
        self.assertEqual(popen.call_args.kwargs["cwd"], "/root")
        process.wait.assert_called_once_with(timeout=30)

    def test_preserves_empty_output(self) -> None:
        _process, start = fake_process(returncode=0)

        with patch("server.subprocess.Popen", side_effect=start):
            result = exec_vps("true")

        self.assertEqual(result["stdout"], "")
        self.assertEqual(result["stderr"], "")
        self.assertEqual(result["exit_code"], 0)

    def test_reports_timeout_and_clamps_to_120_seconds(self) -> None:
        process, start = fake_process(b"partial\n", returncode=-9)
        process.wait.side_effect = [subprocess.TimeoutExpired("sleep", 120), -9]

        with patch("server.subprocess.Popen", side_effect=start):
            result = exec_vps("sleep 999", 999)

        self.assertEqual(result["exit_code"], -9)
        self.assertEqual(result["stdout"], "partial\n")
        self.assertTrue(result["timed_out"])
        process.kill.assert_called_once_with()
        self.assertEqual(process.wait.call_args_list, [call(timeout=120), call()])

    def test_returns_subprocess_exception_as_structured_stderr(self) -> None:
        with patch("server.subprocess.Popen", side_effect=OSError("cannot start")):
            result = exec_vps("example")

        self.assertEqual(
            result,
            {
                "exit_code": None,
                "stdout": "",
                "stderr": "OSError: cannot start",
                "timed_out": False,
                "truncated": False,
            },
        )

    def test_truncates_stdout_and_stderr_independently(self) -> None:
        _process, start = fake_process(
            b"x" * (MAX_OUTPUT_BYTES + 10), b"y" * (MAX_OUTPUT_BYTES + 20)
        )

        with patch("server.subprocess.Popen", side_effect=start):
            result = exec_vps("large-output")

        self.assertEqual(len(result["stdout"].encode()), MAX_OUTPUT_BYTES)
        self.assertEqual(len(result["stderr"].encode()), MAX_OUTPUT_BYTES)
        self.assertTrue(result["truncated"])


if __name__ == "__main__":
    unittest.main()