Aether/tests/services/test_aware_scheduler_quota_skipping.py

from __future__ import annotations

from unittest.mock import MagicMock, patch

from src.services.scheduling.aware_scheduler import CacheAwareScheduler


def _make_key(
    *,
    upstream_metadata: dict,
    allowed_models: list[str] | None = None,
    capabilities: dict[str, bool] | None = None,
) -> MagicMock:
    key = MagicMock()
    key.id = "k1234567890"
    key.allowed_models = allowed_models
    key.capabilities = capabilities or {}
    key.upstream_metadata = upstream_metadata
    return key


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_kiro_quota_remaining_zero_skips(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(upstream_metadata={"kiro": {"remaining": 0.0}})

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="openai:chat",
        model_name="any-model",
        provider_type="kiro",
    )

    assert ok is False
    assert reason == "Kiro 账号配额剩余 0"


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_kiro_quota_remaining_positive_allows(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(upstream_metadata={"kiro": {"remaining": 1.0}})

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="openai:chat",
        model_name="any-model",
        provider_type="kiro",
    )

    assert ok is True
    assert reason is None


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_codex_weekly_quota_exhausted_skips(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(
        upstream_metadata={
            "codex": {
                "primary_used_percent": 100.0,
                "secondary_used_percent": 10.0,
            }
        }
    )

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="openai:cli",
        model_name="any-model",
        provider_type="codex",
    )

    assert ok is False
    assert reason == "Codex 周限额剩余 0%"


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_codex_5h_quota_exhausted_skips(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(
        upstream_metadata={
            "codex": {
                "primary_used_percent": 10.0,
                "secondary_used_percent": 100.0,
            }
        }
    )

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="openai:cli",
        model_name="any-model",
        provider_type="codex",
    )

    assert ok is False
    assert reason == "Codex 5H 限额剩余 0%"


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_codex_ignores_unrelated_metadata_fields(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(
        upstream_metadata={
            "codex": {
                "primary_used_percent": 10.0,
                "secondary_used_percent": 20.0,
                "legacy_marker": "ignore-me",
            }
        }
    )

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="openai:cli",
        model_name="any-model",
        provider_type="codex",
    )

    assert ok is True
    assert reason is None


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_antigravity_model_quota_exhausted_skips(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(
        upstream_metadata={
            "antigravity": {
                "quota_by_model": {
                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
                    "other": {"remaining_fraction": 1.0, "used_percent": 0.0},
                }
            }
        },
    )

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="gemini:chat",
        model_name="ag-model",
        provider_type="antigravity",
    )

    assert ok is False
    assert reason == "Antigravity 模型 ag-model 配额剩余 0%"


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_antigravity_other_model_not_exhausted_allows(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()
    key = _make_key(
        upstream_metadata={
            "antigravity": {
                "quota_by_model": {
                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
                    "other": {"remaining_fraction": 1.0, "used_percent": 0.0},
                }
            }
        },
    )

    ok, reason, _mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="gemini:chat",
        model_name="other",
        provider_type="antigravity",
    )

    assert ok is True
    assert reason is None


@patch("src.services.scheduling.candidate_builder.get_health_monitor")
def test_antigravity_quota_uses_mapping_matched_model(mock_get_health_monitor: MagicMock) -> None:
    mock_get_health_monitor.return_value.get_circuit_breaker_status.return_value = (True, None)
    scheduler = CacheAwareScheduler()

    # Request uses GlobalModel.name, but allowed_models only contains provider-side model id.
    key = _make_key(
        upstream_metadata={
            "antigravity": {
                "quota_by_model": {
                    "ag-model": {"remaining_fraction": 0.0, "used_percent": 100.0},
                }
            }
        },
        allowed_models=["ag-model"],
    )

    ok, reason, mapped = scheduler._candidate_builder._check_key_availability(
        key,
        api_format="gemini:chat",
        model_name="global-model",
        model_mappings=["ag-.*"],
        provider_type="antigravity",
    )

    assert mapped == "ag-model"
    assert ok is False
    assert reason == "Antigravity 模型 ag-model 配额剩余 0%"