Aether/tests/services/test_provider_delete_cleanup.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from src.services.provider import delete_cleanup as cleanup_module
from src.services.provider.delete_cleanup import (
    delete_provider_tree,
    prune_allowed_provider_list,
    prune_allowed_provider_refs,
)


def _query_mock(
    *,
    rows: list[object] | None = None,
    update_count: int | None = None,
    delete_count: int | None = None,
) -> MagicMock:
    query = MagicMock()
    filtered = query.filter.return_value
    if rows is not None:
        filtered.all.return_value = rows
    if update_count is not None:
        filtered.update.return_value = update_count
    if delete_count is not None:
        filtered.delete.return_value = delete_count
    return query


def test_prune_allowed_provider_list_removes_target_and_normalizes_empty() -> None:
    next_allowed, changed = prune_allowed_provider_list(["provider-a", "provider-b"], "provider-a")

    assert changed is True
    assert next_allowed == ["provider-b"]

    next_allowed, changed = prune_allowed_provider_list(["provider-a"], "provider-a")

    assert changed is True
    assert next_allowed == []


def test_prune_allowed_provider_refs_updates_matching_records_only() -> None:
    records = [
        SimpleNamespace(allowed_providers=["provider-a", "provider-b"]),
        SimpleNamespace(allowed_providers=["provider-b"]),
        SimpleNamespace(allowed_providers=None),
    ]

    updated = prune_allowed_provider_refs(records, "provider-a")

    assert updated == 1
    assert records[0].allowed_providers == ["provider-b"]
    assert records[1].allowed_providers == ["provider-b"]
    assert records[2].allowed_providers is None


@pytest.mark.asyncio
async def test_cleanup_deleted_provider_references_cleans_large_fanout_tables(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db = MagicMock()

    user_records = [SimpleNamespace(allowed_providers=["provider-a", "provider-b"])]
    api_key_records = [SimpleNamespace(allowed_providers=["provider-a"])]
    key_cleanup_calls: list[list[str]] = []

    monkeypatch.setattr(
        cleanup_module,
        "cleanup_key_references",
        lambda _db, key_ids: key_cleanup_calls.append(list(key_ids)),
    )

    db.query.side_effect = [
        _query_mock(rows=[("endpoint-1",), ("endpoint-2",)]),
        _query_mock(rows=[("key-1",), ("key-2",)]),
        _query_mock(rows=user_records),
        _query_mock(rows=api_key_records),
        _query_mock(update_count=2),
        _query_mock(update_count=3),
        _query_mock(update_count=4),
        _query_mock(update_count=5),
        _query_mock(update_count=6),
        _query_mock(delete_count=7),
        _query_mock(delete_count=8),
    ]

    stats = cleanup_module.cleanup_deleted_provider_references(db, "provider-a")

    assert stats == {
        "users": 1,
        "api_keys": 1,
        "user_preferences": 2,
        "usage_provider": 3,
        "usage_endpoint": 5,
        "video_tasks_provider": 4,
        "video_tasks_endpoint": 6,
        "request_candidates_provider": 8,
        "request_candidates_endpoint": 7,
    }
    assert user_records[0].allowed_providers == ["provider-b"]
    assert api_key_records[0].allowed_providers == []
    assert key_cleanup_calls == [["key-1", "key-2"]]


def test_delete_provider_tree_deletes_children_before_provider(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db = MagicMock()

    cleanup_mock = MagicMock(
        return_value={
            "users": 1,
            "api_keys": 2,
            "user_preferences": 3,
            "usage_provider": 4,
            "usage_endpoint": 5,
            "video_tasks_provider": 6,
            "video_tasks_endpoint": 7,
            "request_candidates_provider": 8,
            "request_candidates_endpoint": 9,
        }
    )
    monkeypatch.setattr(cleanup_module, "cleanup_deleted_provider_references", cleanup_mock)

    db.query.side_effect = [
        _query_mock(rows=[("endpoint-1",), ("endpoint-2",)]),
        _query_mock(rows=[("key-1",), ("key-2",), ("key-3",)]),
        _query_mock(delete_count=10),
        _query_mock(delete_count=11),
        _query_mock(delete_count=12),
        _query_mock(delete_count=13),
        _query_mock(delete_count=14),
        _query_mock(delete_count=1),
    ]

    stats = delete_provider_tree(db, "provider-a")

    cleanup_mock.assert_called_once_with(
        db,
        "provider-a",
        endpoint_ids=["endpoint-1", "endpoint-2"],
        key_ids=["key-1", "key-2", "key-3"],
    )
    assert stats == {
        "cleanup": cleanup_mock.return_value,
        "deleted": {
            "api_key_mappings": 10,
            "usage_tracking": 11,
            "models": 12,
            "api_keys": 13,
            "endpoints": 14,
            "providers": 1,
        },
        "key_count": 3,
        "endpoint_count": 2,
    }