update agent

agent/tests/test_apt_components.py

@@ -0,0 +1,205 @@
+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()