InstallerRobot/agent/tests/test_apt_components.py

from __future__ import annotations

import subprocess
import unittest
from types import SimpleNamespace
from unittest.mock import patch

from app.core.installer import APT_NONINTERACTIVE_ENV, AptInstaller
from app.core.manifest_validator import ManifestValidator
from app.core.task_runner import TaskRunner


def apt_manifest(component: dict | None = None) -> dict:
    return {
        "schemaVersion": "1.0",
        "appId": "postgresql",
        "appName": "PostgreSQL",
        "version": "16",
        "components": [
            component
            or {
                "componentId": "postgresql",
                "type": "apt",
                "packageName": "postgresql",
                "version": "16",
            }
        ],
    }


class FakeCommandRunner:
    def __init__(self) -> None:
        self.calls: list[tuple[list[str], int | None, dict[str, str] | None]] = []

    def run(
        self,
        command: list[str],
        timeout: int | None = None,
        env: dict[str, str] | None = None,
    ) -> subprocess.CompletedProcess[str]:
        self.calls.append((command, timeout, env))
        return subprocess.CompletedProcess(command, 0, stdout="accepting connections\n", stderr="")


class AptManifestTests(unittest.TestCase):
    def test_postgresql_is_accepted_when_allowlisted(self) -> None:
        validator_settings = SimpleNamespace(allowed_apt_packages=["postgresql"])
        with patch("app.core.manifest_validator.settings", validator_settings):
            manifest = ManifestValidator().validate(apt_manifest())

        self.assertEqual(manifest["components"][0]["type"], "apt")
        self.assertEqual(manifest["components"][0]["packageName"], "postgresql")

    def test_non_allowlisted_package_is_rejected(self) -> None:
        validator_settings = SimpleNamespace(allowed_apt_packages=["postgresql"])
        with patch("app.core.manifest_validator.settings", validator_settings):
            with self.assertRaisesRegex(ValueError, "APT package is not allowed"):
                ManifestValidator().validate(
                    apt_manifest(
                        {
                            "componentId": "curl",
                            "type": "apt",
                            "packageName": "curl",
                        }
                    )
                )

    def test_manifest_cannot_supply_a_shell_command(self) -> None:
        validator_settings = SimpleNamespace(allowed_apt_packages=["postgresql"])
        with patch("app.core.manifest_validator.settings", validator_settings):
            with self.assertRaises(ValueError):
                ManifestValidator().validate(
                    apt_manifest(
                        {
                            "componentId": "postgresql",
                            "type": "apt",
                            "packageName": "postgresql",
                            "command": "rm -rf /",
                        }
                    )
                )


class AptInstallerTests(unittest.TestCase):
    def test_apt_commands_are_fixed_argument_lists(self) -> None:
        runner = FakeCommandRunner()
        installer = AptInstaller(runner)

        installer.update_package_index()
        installer.install_package("postgresql")
        installer.wait_for_postgresql(attempts=1, delay_seconds=0)

        self.assertEqual(runner.calls[0][0], ["apt-get", "update"])
        self.assertEqual(
            runner.calls[1][0],
            [
                "apt-get",
                "-o",
                "Dpkg::Use-Pty=0",
                "-o",
                "Dpkg::Options::=--force-confdef",
                "-o",
                "Dpkg::Options::=--force-confold",
                "install",
                "--yes",
                "postgresql",
            ],
        )
        self.assertEqual(runner.calls[1][2], APT_NONINTERACTIVE_ENV)
        self.assertEqual(runner.calls[2][0], ["pg_isready", "--timeout=5"])


class FakeRepository:
    def __init__(self) -> None:
        self.component_updates: list[dict] = []
        self.logs: list[tuple[str, str]] = []
        self.installed_component: dict | None = None

    def update_task_component(self, task_id: str, component_id: str, **fields: object) -> None:
        self.component_updates.append(dict(fields))

    def add_log(self, task_id: str, level: str, message: str) -> None:
        self.logs.append((level, message))

    def upsert_installed_component(self, app_id: str, component: dict) -> None:
        self.installed_component = component


class FakeAptInstaller:
    actions: list[str] = []

    def __init__(self, command_runner: object) -> None:
        self.command_runner = command_runner

    def update_package_index(self) -> None:
        self.actions.append("update")

    def install_package(self, package_name: str) -> None:
        self.actions.append(f"install:{package_name}")

    def get_package_version(self, package_name: str) -> str:
        self.actions.append(f"version:{package_name}")
        return "16+257build1"

    def wait_for_postgresql(self) -> None:
        self.actions.append("pg_isready")


class FakeServiceManager:
    actions: list[str] = []

    def __init__(self, command_runner: object) -> None:
        self.command_runner = command_runner

    def enable_service(self, service_name: str) -> None:
        self.actions.append(f"enable:{service_name}")

    def start_service(self, service_name: str) -> None:
        self.actions.append(f"start:{service_name}")

    def assert_service_active(self, service_name: str) -> None:
        self.actions.append(f"active:{service_name}")


class AptTaskRunnerTests(unittest.TestCase):
    def test_postgresql_install_verifies_service_and_readiness(self) -> None:
        repository = FakeRepository()
        FakeAptInstaller.actions = []
        FakeServiceManager.actions = []

        with (
            patch("app.core.task_runner.CommandRunner", return_value=object()),
            patch("app.core.task_runner.AptInstaller", FakeAptInstaller),
            patch("app.core.task_runner.ServiceManager", FakeServiceManager),
        ):
            TaskRunner(repository)._install_apt_component(
                "task-1",
                "postgresql-app",
                {
                    "componentId": "postgresql",
                    "type": "apt",
                    "packageName": "postgresql",
                    "version": "16",
                },
            )

        self.assertEqual(
            FakeAptInstaller.actions,
            ["update", "install:postgresql", "version:postgresql", "pg_isready"],
        )
        self.assertEqual(
            FakeServiceManager.actions,
            [
                "enable:postgresql.service",
                "start:postgresql.service",
                "active:postgresql.service",
            ],
        )
        self.assertEqual(repository.installed_component["version"], "16+257build1")
        self.assertEqual(repository.installed_component["serviceName"], "postgresql.service")
        self.assertEqual(repository.component_updates[-1]["current_step"], "checking PostgreSQL readiness")


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