dashboard-nanobot/backend/tests/test_docker_manager.py

import sys
import tempfile
import types
import unittest
from unittest.mock import MagicMock

docker_stub = types.ModuleType("docker")
docker_stub.errors = types.SimpleNamespace(
    ImageNotFound=type("ImageNotFound", (Exception,), {}),
    NotFound=type("NotFound", (Exception,), {}),
)
sys.modules.setdefault("docker", docker_stub)

from core.docker_manager import BotDockerManager


class BotDockerManagerTests(unittest.TestCase):
    def setUp(self) -> None:
        self._tmpdir = tempfile.TemporaryDirectory()

    def tearDown(self) -> None:
        self._tmpdir.cleanup()

    def _make_manager(self) -> BotDockerManager:
        manager = BotDockerManager.__new__(BotDockerManager)
        manager.client = MagicMock()
        manager.host_data_root = self._tmpdir.name
        manager.base_image = "nanobot-base"
        manager.network_name = ""
        manager.active_monitors = {}
        manager._last_delivery_error = {}
        manager._storage_limit_supported = True
        manager._storage_limit_warning_emitted = False
        return manager

    @staticmethod
    def _build_container(
        *,
        status: str,
        image: str,
        nano_cpus: int,
        memory_bytes: int,
        storage_opt_size: str,
        source_mount: str,
        network_name: str,
        bootstrap_label: str | None = "env-json-v1",
    ) -> MagicMock:
        container = MagicMock()
        container.status = status
        container.reload = MagicMock()
        container.start = MagicMock()
        container.stop = MagicMock()
        container.remove = MagicMock()
        container.attrs = {
            "Config": {
                "Image": image,
                "Labels": (
                    {"dashboard.runtime_bootstrap": bootstrap_label}
                    if bootstrap_label is not None
                    else {}
                ),
            },
            "HostConfig": {
                "NanoCpus": nano_cpus,
                "Memory": memory_bytes,
                "StorageOpt": {"size": storage_opt_size},
            },
            "Mounts": [
                {
                    "Source": source_mount,
                    "Destination": "/root/.nanobot",
                    "RW": True,
                }
            ],
            "NetworkSettings": {
                "Networks": {network_name: {"IPAddress": "172.18.0.2"}},
            },
        }
        return container

    def test_stop_bot_keeps_container_by_default(self) -> None:
        manager = self._make_manager()
        container = MagicMock()
        container.status = "running"
        container.reload = MagicMock()
        container.stop = MagicMock()
        container.remove = MagicMock()
        manager.client.containers.get.return_value = container

        result = manager.stop_bot("demo")

        self.assertTrue(result)
        container.stop.assert_called_once_with(timeout=5)
        container.remove.assert_not_called()

    def test_stop_bot_remove_true_deletes_container(self) -> None:
        manager = self._make_manager()
        container = MagicMock()
        container.status = "exited"
        container.reload = MagicMock()
        container.stop = MagicMock()
        container.remove = MagicMock()
        manager.client.containers.get.return_value = container

        result = manager.stop_bot("demo", remove=True)

        self.assertTrue(result)
        container.stop.assert_not_called()
        container.remove.assert_called_once_with()

    def test_start_bot_reuses_compatible_stopped_container(self) -> None:
        manager = self._make_manager()
        image_tag = "nanobot-base:v1"
        workspace_mount = f"{self._tmpdir.name}/demo/.nanobot"
        container = self._build_container(
            status="exited",
            image=image_tag,
            nano_cpus=1_000_000_000,
            memory_bytes=1024 * 1024 * 1024,
            storage_opt_size="10G",
            source_mount=workspace_mount,
            network_name="bridge",
        )
        manager.client.images.get.return_value = MagicMock()
        manager.client.containers.get.return_value = container

        result = manager.start_bot(
            "demo",
            image_tag=image_tag,
            env_vars={"TZ": "UTC", "API_KEY": "updated-secret"},
            cpu_cores=1.0,
            memory_mb=1024,
            storage_gb=10,
        )

        self.assertTrue(result)
        container.start.assert_called_once_with()
        container.remove.assert_not_called()
        manager.client.containers.run.assert_not_called()

    def test_start_bot_recreates_incompatible_stopped_container(self) -> None:
        manager = self._make_manager()
        image_tag = "nanobot-base:v1"
        workspace_mount = f"{self._tmpdir.name}/demo/.nanobot"
        container = self._build_container(
            status="exited",
            image="nanobot-base:old",
            nano_cpus=1_000_000_000,
            memory_bytes=1024 * 1024 * 1024,
            storage_opt_size="10G",
            source_mount=workspace_mount,
            network_name="bridge",
        )
        manager.client.images.get.return_value = MagicMock()
        manager.client.containers.get.return_value = container
        manager._run_container_with_storage_fallback = MagicMock(return_value=MagicMock())

        result = manager.start_bot(
            "demo",
            image_tag=image_tag,
            env_vars={"TZ": "Asia/Shanghai"},
            cpu_cores=1.0,
            memory_mb=1024,
            storage_gb=10,
        )

        self.assertTrue(result)
        container.start.assert_not_called()
        container.remove.assert_called_once_with(force=True)
        manager._run_container_with_storage_fallback.assert_called_once()

    def test_start_bot_recreates_container_without_new_entrypoint(self) -> None:
        manager = self._make_manager()
        image_tag = "nanobot-base:v1"
        workspace_mount = f"{self._tmpdir.name}/demo/.nanobot"
        container = self._build_container(
            status="exited",
            image=image_tag,
            nano_cpus=1_000_000_000,
            memory_bytes=1024 * 1024 * 1024,
            storage_opt_size="10G",
            source_mount=workspace_mount,
            network_name="bridge",
            bootstrap_label=None,
        )
        manager.client.images.get.return_value = MagicMock()
        manager.client.containers.get.return_value = container
        manager._run_container_with_storage_fallback = MagicMock(return_value=MagicMock())

        result = manager.start_bot(
            "demo",
            image_tag=image_tag,
            env_vars={"TZ": "Asia/Shanghai"},
            cpu_cores=1.0,
            memory_mb=1024,
            storage_gb=10,
        )

        self.assertTrue(result)
        container.start.assert_not_called()
        container.remove.assert_called_once_with(force=True)
        manager._run_container_with_storage_fallback.assert_called_once()


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