devplacepy/tests/unit/services/devii/container/controller.py

# retoor <retoor@molodetz.nl>

import json
from datetime import datetime, timedelta, timezone

import pytest

from devplacepy.database import get_table, invalidate_admins_cache
from devplacepy.services.devii.container.controller import ContainerController
from devplacepy.services.devii.errors import ToolInputError
from devplacepy.utils import generate_uid

from tests.conftest import run_async


def _make_user_at(role, created_at):
    uid = generate_uid()
    username = f"cc_{uid[:8]}"
    get_table("users").insert(
        {
            "uid": uid,
            "username": username,
            "email": f"{username}@t.dev",
            "role": role,
            "created_at": created_at.isoformat(),
        }
    )
    invalidate_admins_cache()
    return get_table("users").find_one(uid=uid)


def _demote_existing_admins():
    existing = [r["uid"] for r in get_table("users").find(role="Admin")]
    for uid in existing:
        get_table("users").update({"uid": uid, "role": "Member"}, ["uid"])
    invalidate_admins_cache()
    return existing


def _restore_admins(uids):
    for uid in uids:
        get_table("users").update({"uid": uid, "role": "Admin"}, ["uid"])
    invalidate_admins_cache()


def _purge(*rows, project_uids=(), instance_uids=()):
    for row in rows:
        if row:
            get_table("users").delete(uid=row["uid"])
    for uid in project_uids:
        get_table("projects").delete(uid=uid)
    for uid in instance_uids:
        get_table("instances").delete(uid=uid)


def _make_project(owner_uid, is_private):
    uid = generate_uid()
    slug = f"cc-{uid[:10]}"
    get_table("projects").insert(
        {
            "uid": uid,
            "slug": slug,
            "title": "Controller Test Project",
            "user_uid": owner_uid,
            "is_private": 1 if is_private else 0,
            "deleted_at": None,
            "deleted_by": None,
        }
    )
    return get_table("projects").find_one(uid=uid)


def _make_instance(project_uid, created_by):
    uid = generate_uid()
    get_table("instances").insert(
        {
            "uid": uid,
            "slug": f"cci-{uid[:10]}",
            "name": "controller-test",
            "project_uid": project_uid,
            "created_by": created_by,
            "owner_uid": "",
            "status": "created",
            "container_id": "",
            "deleted_at": None,
            "deleted_by": None,
        }
    )
    return uid


def _controller(owner_uid):
    return ContainerController(client=None, owner_id=owner_uid)


def test_project_lookup_denied_for_non_viewer_on_private_project(local_db):
    restore = _demote_existing_admins()
    owner_admin = other_admin = None
    project = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=True)
        controller = _controller(other_admin["uid"])
        with pytest.raises(ToolInputError):
            run_async(
                controller.dispatch(
                    "container_list_instances", {"project_slug": project["slug"]}
                )
            )
    finally:
        _purge(
            owner_admin,
            other_admin,
            project_uids=[project["uid"]] if project else [],
        )
        _restore_admins(restore)


def test_project_lookup_allowed_for_owner(local_db):
    restore = _demote_existing_admins()
    owner_admin = None
    project = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        project = _make_project(owner_admin["uid"], is_private=True)
        controller = _controller(owner_admin["uid"])
        out = json.loads(
            run_async(
                controller.dispatch(
                    "container_list_instances", {"project_slug": project["slug"]}
                )
            )
        )
        assert out["instances"] == []
    finally:
        _purge(owner_admin, project_uids=[project["uid"]] if project else [])
        _restore_admins(restore)


def test_project_lookup_allowed_for_primary_admin_on_others_private_project(local_db):
    restore = _demote_existing_admins()
    primary = other_admin = None
    project = None
    try:
        base = datetime.now(timezone.utc)
        primary = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(other_admin["uid"], is_private=True)
        controller = _controller(primary["uid"])
        out = json.loads(
            run_async(
                controller.dispatch(
                    "container_list_instances", {"project_slug": project["slug"]}
                )
            )
        )
        assert out["instances"] == []
    finally:
        _purge(
            primary,
            other_admin,
            project_uids=[project["uid"]] if project else [],
        )
        _restore_admins(restore)


def test_instance_action_denied_for_non_owner_admin_on_public_project(local_db):
    restore = _demote_existing_admins()
    owner_admin = other_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(other_admin["uid"])
        with pytest.raises(ToolInputError):
            run_async(
                controller.dispatch(
                    "container_instance_action",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "action": "stop",
                    },
                )
            )
    finally:
        _purge(
            owner_admin,
            other_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_instance_action_allowed_for_owner(local_db):
    restore = _demote_existing_admins()
    owner_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(owner_admin["uid"])
        out = json.loads(
            run_async(
                controller.dispatch(
                    "container_instance_action",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "action": "stop",
                    },
                )
            )
        )
        assert out["status"] == "ok"
    finally:
        _purge(
            owner_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_instance_action_allowed_for_primary_admin_on_others_instance(local_db):
    restore = _demote_existing_admins()
    primary = owner_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        primary = _make_user_at("Admin", base)
        owner_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=True)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(primary["uid"])
        out = json.loads(
            run_async(
                controller.dispatch(
                    "container_instance_action",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "action": "stop",
                    },
                )
            )
        )
        assert out["status"] == "ok"
    finally:
        _purge(
            primary,
            owner_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_configure_instance_denied_for_non_owner(local_db):
    restore = _demote_existing_admins()
    owner_admin = other_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(other_admin["uid"])
        with pytest.raises(ToolInputError):
            run_async(
                controller.dispatch(
                    "container_configure_instance",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "restart_policy": "always",
                    },
                )
            )
    finally:
        _purge(
            owner_admin,
            other_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_exec_denied_for_non_owner_before_running_check(local_db):
    restore = _demote_existing_admins()
    owner_admin = other_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(other_admin["uid"])
        with pytest.raises(ToolInputError) as exc:
            run_async(
                controller.dispatch(
                    "container_exec",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "command": "ls",
                    },
                )
            )
        message = str(exc.value).lower()
        assert "owner" in message or "primary" in message
    finally:
        _purge(
            owner_admin,
            other_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_exec_allowed_for_owner_reaches_running_check(local_db):
    restore = _demote_existing_admins()
    owner_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(owner_admin["uid"])
        with pytest.raises(ToolInputError) as exc:
            run_async(
                controller.dispatch(
                    "container_exec",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "command": "ls",
                    },
                )
            )
        assert "not running" in str(exc.value).lower()
    finally:
        _purge(
            owner_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_schedule_denied_for_non_owner(local_db):
    restore = _demote_existing_admins()
    owner_admin = other_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        other_admin = _make_user_at("Admin", base + timedelta(seconds=1))
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(other_admin["uid"])
        with pytest.raises(ToolInputError):
            run_async(
                controller.dispatch(
                    "container_schedule",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "action": "start",
                        "kind": "cron",
                        "cron": "0 3 * * *",
                    },
                )
            )
    finally:
        _purge(
            owner_admin,
            other_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_schedule_allowed_for_owner(local_db):
    restore = _demote_existing_admins()
    owner_admin = None
    project = None
    inst_uid = None
    try:
        base = datetime.now(timezone.utc)
        owner_admin = _make_user_at("Admin", base)
        project = _make_project(owner_admin["uid"], is_private=False)
        inst_uid = _make_instance(project["uid"], owner_admin["uid"])
        controller = _controller(owner_admin["uid"])
        out = json.loads(
            run_async(
                controller.dispatch(
                    "container_schedule",
                    {
                        "project_slug": project["slug"],
                        "instance": inst_uid,
                        "action": "start",
                        "kind": "cron",
                        "cron": "0 3 * * *",
                    },
                )
            )
        )
        assert out["schedule"]["action"] == "start"
    finally:
        _purge(
            owner_admin,
            project_uids=[project["uid"]] if project else [],
            instance_uids=[inst_uid] if inst_uid else [],
        )
        _restore_admins(restore)


def test_unknown_project_slug_raises(local_db):
    restore = _demote_existing_admins()
    owner_admin = None
    try:
        owner_admin = _make_user_at("Admin", datetime.now(timezone.utc))
        controller = _controller(owner_admin["uid"])
        with pytest.raises(ToolInputError):
            run_async(
                controller.dispatch(
                    "container_list_instances", {"project_slug": "no-such-slug"}
                )
            )
    finally:
        _purge(owner_admin)
        _restore_admins(restore)