Aether/tests/services/test_global_model_auto_disassociate.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock

from src.models.database import Model, Provider, ProviderAPIKey
from src.services.model.global_model import GlobalModelService


def test_auto_disassociate_short_circuits_when_unlimited_key_exists() -> None:
    db = MagicMock()
    provider_query = MagicMock()
    provider_query.filter.return_value.first.return_value = SimpleNamespace(name="Provider A")

    unlimited_query = MagicMock()
    unlimited_query.filter.return_value.limit.return_value.first.return_value = object()

    def _query(*entities: object) -> MagicMock:
        entity = entities[0]
        if entity is Provider:
            return provider_query
        if entity is ProviderAPIKey.id:
            return unlimited_query
        if entity is Model:
            raise AssertionError("model query should not run when unlimited key exists")
        raise AssertionError(f"unexpected query: {entities}")

    db.query.side_effect = _query

    result = GlobalModelService.auto_disassociate_provider_by_key_whitelist(db, "provider-1")

    assert result == {"success": [], "errors": []}
    db.delete.assert_not_called()
    db.commit.assert_not_called()


def test_auto_disassociate_deletes_unmatched_auto_associated_models(
    monkeypatch,
) -> None:
    db = MagicMock()
    provider_query = MagicMock()
    provider_query.filter.return_value.first.return_value = SimpleNamespace(name="Provider B")

    unlimited_query = MagicMock()
    unlimited_query.filter.return_value.limit.return_value.first.return_value = None

    allowed_models_query = MagicMock()
    # db.query(ProviderAPIKey.allowed_models).all() returns list of tuples
    allowed_models_query.filter.return_value.all.return_value = [
        (["gpt-4o"],),
        ([],),
    ]

    model = SimpleNamespace(
        id="model-1",
        global_model=SimpleNamespace(
            id="gm-1",
            name="claude-sonnet",
            config={"model_mappings": ["claude-*"]},
        ),
    )
    models_query = MagicMock()
    models_query.options.return_value.filter.return_value.all.return_value = [model]

    def _query(*entities: object) -> MagicMock:
        entity = entities[0]
        if entity is Provider:
            return provider_query
        if entity is ProviderAPIKey.id:
            return unlimited_query
        if entity is ProviderAPIKey.allowed_models:
            return allowed_models_query
        if entity is Model:
            return models_query
        raise AssertionError(f"unexpected query: {entities}")

    db.query.side_effect = _query
    monkeypatch.setattr(
        "src.core.model_permissions.match_model_with_pattern",
        lambda pattern, allowed_model: pattern == allowed_model,
    )

    result = GlobalModelService.auto_disassociate_provider_by_key_whitelist(db, "provider-2")

    assert result["errors"] == []
    assert result["success"] == [
        {
            "model_id": "model-1",
            "global_model_id": "gm-1",
            "global_model_name": "claude-sonnet",
        }
    ]
    db.delete.assert_called_once_with(model)
    db.commit.assert_called_once()