Octop/tests/unit/test_storage_spec.py

"""Tests for storage backend spec resolution."""

from __future__ import annotations

from pathlib import Path

import pytest

from octop.infra.backend.adapter import row_to_backend_spec, storage_backend_kind_agent_resolvable
from octop.infra.backend.resolver import (
    backend_spec_supports_execution,
    collect_named_storage_backend_refs,
    find_agents_using_storage_backend,
    resolve_agent_backend_spec,
    rewrite_named_storage_backend_refs,
)
from octop.infra.db.migrate import run_migrations
from octop.infra.db.pool import SqlitePool
from octop.infra.db.repos.backends import BackendRepo, BackendRow


@pytest.fixture
def repo(tmp_path: Path) -> BackendRepo:
    db = SqlitePool(tmp_path / "octop.db")
    run_migrations(db)
    return BackendRepo(db)


def test_resolve_named_backend(repo: BackendRepo) -> None:
    repo.create(
        name="my-cos",
        kind="cos",
        access_key="AKID",
        secret_key="SECRET",
        bucket="b-125",
        region="ap-guangzhou",
    )
    resolved = resolve_agent_backend_spec(
        {"type": "named", "name": "my-cos"},
        repo=repo,
    )
    assert resolved == {
        "type": "cos",
        "bucket": "b-125",
        "region": "ap-guangzhou",
        "secret_id": "AKID",
        "secret_key": "SECRET",
    }


def test_resolve_composite_with_named(repo: BackendRepo) -> None:
    repo.create(
        name="archive",
        kind="filesystem",
        bucket="/data/archive",
    )
    spec = {
        "type": "composite",
        "default": {"type": "filesystem", "virtual_mode": True},
        "routes": {"/archive": {"type": "named", "name": "archive"}},
    }
    resolved = resolve_agent_backend_spec(spec, repo=repo)
    archive_row = repo.get_by_name("archive")
    assert archive_row is not None
    assert resolved["type"] == "composite"
    assert resolved["routes"]["/archive"] == row_to_backend_spec(archive_row)


def test_backend_spec_supports_execution() -> None:
    assert backend_spec_supports_execution({"type": "local_shell", "virtual_mode": True})
    assert backend_spec_supports_execution({"type": "docker", "image": "python:3.12-slim"})
    assert backend_spec_supports_execution("docker")
    assert backend_spec_supports_execution({"type": "opensandbox", "image": "python:3.12"})
    assert backend_spec_supports_execution("opensandbox")
    assert not backend_spec_supports_execution({"type": "filesystem", "virtual_mode": True})
    assert not backend_spec_supports_execution({"type": "state"})
    assert backend_spec_supports_execution(
        {
            "type": "composite",
            "default": {"type": "filesystem", "virtual_mode": True},
            "routes": {"/shell": {"type": "local_shell", "virtual_mode": True}},
        }
    )


def test_storage_backend_kind_agent_resolvable() -> None:
    assert storage_backend_kind_agent_resolvable("cos")
    assert storage_backend_kind_agent_resolvable("filesystem")
    assert storage_backend_kind_agent_resolvable("docker")
    assert storage_backend_kind_agent_resolvable("opensandbox")
    assert storage_backend_kind_agent_resolvable("postgres")


def test_row_to_backend_spec_postgres_discrete_fields() -> None:
    row = BackendRow(
        id=1,
        name="pg",
        kind="postgres",
        endpoint="db.internal:5433",
        access_key="octop",
        secret_key="p@ss/word#1",
        bucket="octop",
        region="agent",
        config_json='{"table": "files"}',
        note=None,
        enabled=1,
        created_at=0,
        updated_at=0,
    )
    spec = row_to_backend_spec(row)
    assert spec == {
        "type": "postgres",
        "host": "db.internal",
        "port": 5433,
        "user": "octop",
        "password": "p@ss/word#1",
        "database": "octop",
        "schema": "agent",
        "table": "files",
    }
    assert "connection_string" not in spec


def test_row_to_backend_spec_postgres_parses_uri() -> None:
    row = BackendRow(
        id=1,
        name="pg",
        kind="postgres",
        endpoint=None,
        access_key=None,
        secret_key=None,
        bucket=None,
        region=None,
        config_json='{"connection_string": "postgresql://alice:s3cret@pg.example:5432/app?sslmode=require"}',
        note=None,
        enabled=1,
        created_at=0,
        updated_at=0,
    )
    spec = row_to_backend_spec(row)
    assert spec is not None
    assert spec["host"] == "pg.example"
    assert spec["port"] == 5432
    assert spec["user"] == "alice"
    assert spec["password"] == "s3cret"
    assert spec["database"] == "app"
    assert spec["sslmode"] == "require"
    assert "connection_string" not in spec


def test_resolve_inline_postgres_drops_connection_string() -> None:
    resolved = resolve_agent_backend_spec(
        {"type": "postgres", "connection_string": "postgresql://u:p@localhost/db"},
        repo=None,
    )
    assert resolved["type"] == "postgres"
    assert resolved["host"] == "localhost"
    assert resolved["user"] == "u"
    assert resolved["password"] == "p"
    assert resolved["database"] == "db"
    assert "connection_string" not in resolved


def test_row_to_backend_spec_opensandbox() -> None:
    from octop.infra.backend.adapter import storage_spec_previewable

    row = BackendRow(
        id=1,
        name="osb",
        kind="opensandbox",
        endpoint="127.0.0.1:8080",
        access_key=None,
        secret_key="sk-test",
        bucket="python:3.12",
        region="https",
        config_json='{"timeout": 120, "use_server_proxy": true}',
        note=None,
        enabled=1,
        created_at=0,
        updated_at=0,
    )
    assert row_to_backend_spec(row) == {
        "type": "opensandbox",
        "image": "python:3.12",
        "api_key": "sk-test",
        "domain": "127.0.0.1:8080",
        "protocol": "https",
        "timeout": 120,
        "use_server_proxy": True,
    }
    assert storage_spec_previewable(row_to_backend_spec(row)) is False


def test_row_to_backend_spec_docker() -> None:
    row = BackendRow(
        id=1,
        name="sandbox",
        kind="docker",
        endpoint=None,
        access_key=None,
        secret_key=None,
        bucket="python:3.12-slim",
        region=None,
        config_json='{"allow_network": false, "memory": "512m", "sandbox_scope": "user", "sandbox_prefix": "octop_sandbox"}',
        note=None,
        enabled=1,
        created_at=0,
        updated_at=0,
    )
    assert row_to_backend_spec(row) == {
        "type": "docker",
        "image": "python:3.12-slim",
        "allow_network": False,
        "memory": "512m",
        "sandbox_scope": "user",
        "sandbox_prefix": "octop_sandbox",
    }


def test_enrich_docker_backend_spec_defaults() -> None:
    from octop.infra.backend.docker_spec import (
        docker_spec_previewable,
        enrich_docker_backend_spec,
    )

    agent = enrich_docker_backend_spec(
        {"type": "docker", "image": "python:3.12-slim"},
        agent_id="ZE6GR2",
    )
    assert agent["sandbox_prefix"] == "octop_sandbox"
    assert agent["sandbox_scope"] == "agent"
    assert agent["agent_id"] == "ZE6GR2"
    assert docker_spec_previewable(agent) is False

    user = enrich_docker_backend_spec(
        {"type": "docker", "image": "python:3.12-slim", "sandbox_scope": "user"},
        agent_id="ZE6GR2",
        username="alice",
    )
    assert user["username"] == "alice"
    assert docker_spec_previewable(user) is False

    fixed = enrich_docker_backend_spec(
        {
            "type": "docker",
            "image": "python:3.12-slim",
            "sandbox_scope": "fixed",
            "sandbox_id": "shared1",
        },
        agent_id="ZE6GR2",
    )
    assert docker_spec_previewable(fixed) is True
    assert "agent_id" not in fixed or fixed.get("agent_id") == "ZE6GR2"


def test_inject_docker_global_environment_sets_file() -> None:
    from octop.infra.backend.docker_spec import inject_docker_global_environment

    spec = inject_docker_global_environment(
        {"type": "docker", "environment": {"FOO": "from-spec"}},
        "/data/octop/env",
    )
    assert spec["environment_file"] == "/data/octop/env"
    assert spec["environment"] == {"FOO": "from-spec"}
    local = inject_docker_global_environment({"type": "local_shell"}, "/data/octop/env")
    assert "environment_file" not in local


def test_collect_named_storage_backend_refs() -> None:
    assert collect_named_storage_backend_refs({"type": "named", "name": "my-cos"}) == {"my-cos"}
    assert collect_named_storage_backend_refs(
        {
            "type": "composite",
            "default": {"type": "filesystem", "virtual_mode": True},
            "routes": {
                "/data": {"type": "named", "name": "archive"},
                "/x": {"type": "local_shell", "virtual_mode": True},
            },
        }
    ) == {"archive"}
    assert collect_named_storage_backend_refs(None) == set()


def test_find_agents_using_storage_backend() -> None:
    class _Row:
        def __init__(self, agent_id: str, name: str) -> None:
            self.agent_id = agent_id
            self.name = name

    class _Repo:
        def list_all(self) -> list[_Row]:
            return [
                _Row("A1", "Alice Bot"),
                _Row("B2", "Bob Bot"),
            ]

    configs = {
        "A1": {"backend": {"type": "named", "name": "my-cos"}},
        "B2": {"backend": {"type": "filesystem", "virtual_mode": True}},
    }

    refs = find_agents_using_storage_backend(
        agent_repo=_Repo(),
        get_config=lambda aid: configs.get(aid, {}),
        backend_name="my-cos",
    )
    assert refs == [{"agent_id": "A1", "name": "Alice Bot"}]


def test_rewrite_named_storage_backend_refs() -> None:
    assert rewrite_named_storage_backend_refs(
        {"type": "named", "name": "old"},
        old="old",
        new="new",
    ) == {"type": "named", "name": "new"}
    assert rewrite_named_storage_backend_refs("named:old", old="old", new="new") == "named:new"
    assert rewrite_named_storage_backend_refs(
        {
            "type": "composite",
            "default": {"type": "named", "name": "old"},
            "routes": {"/data": {"type": "named", "name": "old"}},
        },
        old="old",
        new="new",
    ) == {
        "type": "composite",
        "default": {"type": "named", "name": "new"},
        "routes": {"/data": {"type": "named", "name": "new"}},
    }
    assert rewrite_named_storage_backend_refs(
        {"type": "named", "name": "keep"},
        old="old",
        new="new",
    ) == {"type": "named", "name": "keep"}