Aether/tests/services/test_startup_toggles.py

from __future__ import annotations

import asyncio
import inspect
from collections.abc import Callable
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock

import pytest

import src.main as main_module
import src.services.system.maintenance_scheduler as maintenance_scheduler_module
from src.config.settings import config
from src.services.system.maintenance_scheduler import MaintenanceScheduler


@pytest.mark.asyncio
async def test_maintenance_scheduler_start_skips_startup_task_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(config, "maintenance_startup_tasks_enabled", False)

    scheduler = MaintenanceScheduler()

    created = False

    def fake_create_task(coro):  # type: ignore[no-untyped-def]
        nonlocal created
        created = True
        if inspect.iscoroutine(coro):
            coro.close()
        return object()

    monkeypatch.setattr(maintenance_scheduler_module.asyncio, "create_task", fake_create_task)
    monkeypatch.setattr(scheduler, "_get_checkin_time", lambda: (1, 5))
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "get_scheduler",
        lambda: SimpleNamespace(
            add_cron_job=lambda *args, **kwargs: None,
            add_interval_job=lambda *args, **kwargs: None,
        ),
    )

    await scheduler.start()

    assert created is False


@pytest.mark.asyncio
async def test_maintenance_scheduler_start_skips_startup_task_when_no_python_owner(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(config, "maintenance_startup_tasks_enabled", True)
    monkeypatch.setattr(config, "pending_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", False)

    scheduler = MaintenanceScheduler()
    safe_create_task = MagicMock()

    monkeypatch.setattr("src.utils.async_utils.safe_create_task", safe_create_task)
    monkeypatch.setattr(scheduler, "_get_checkin_time", lambda: (1, 5))
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "get_scheduler",
        lambda: SimpleNamespace(
            add_cron_job=lambda *args, **kwargs: None,
            add_interval_job=lambda *args, **kwargs: None,
        ),
    )

    await scheduler.start()

    safe_create_task.assert_not_called()
    assert scheduler._startup_task is None


@pytest.mark.asyncio
async def test_maintenance_scheduler_stop_cancels_startup_task_and_removes_jobs(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()
    scheduler.running = True
    scheduler._startup_task = asyncio.create_task(asyncio.sleep(3600))

    expected_job_ids = [
        "stats_aggregation",
        "stats_hourly_aggregation",
        "wallet_daily_usage_aggregation",
        "usage_cleanup",
        "pool_monitor",
        "http_client_idle_cleanup",
        "pending_cleanup",
        "audit_cleanup",
        "gemini_file_mapping_cleanup",
        "candidate_cleanup",
        "db_maintenance",
        "antigravity_ua_refresh",
        scheduler.CHECKIN_JOB_ID,
    ]
    scheduler._registered_job_ids = list(expected_job_ids)
    removed_jobs: list[str] = []

    monkeypatch.setattr(
        maintenance_scheduler_module,
        "get_scheduler",
        lambda: SimpleNamespace(remove_job=lambda job_id: removed_jobs.append(job_id)),
    )

    await scheduler.stop()

    assert scheduler.running is False
    assert scheduler._startup_task is None
    assert set(removed_jobs) == set(expected_job_ids)
    assert scheduler._registered_job_ids == []


@pytest.mark.asyncio
async def test_maintenance_scheduler_start_skips_rust_owned_maintenance_jobs_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(config, "maintenance_startup_tasks_enabled", False)
    monkeypatch.setattr(config, "audit_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "db_maintenance_python_enabled", False)
    monkeypatch.setattr(config, "gemini_file_mapping_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "http_client_idle_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "pending_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "pool_monitor_python_enabled", False)
    monkeypatch.setattr(config, "provider_checkin_python_enabled", False)
    monkeypatch.setattr(config, "request_candidate_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "stats_aggregation_python_enabled", False)
    monkeypatch.setattr(config, "stats_hourly_aggregation_python_enabled", False)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", False)
    monkeypatch.setattr(config, "usage_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "wallet_daily_usage_aggregation_python_enabled", False)

    scheduler = MaintenanceScheduler()
    cron_job_ids: list[str] = []
    interval_job_ids: list[str] = []

    monkeypatch.setattr(scheduler, "_get_checkin_time", lambda: (1, 5))
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "get_scheduler",
        lambda: SimpleNamespace(
            add_cron_job=lambda *, job_id, **kwargs: cron_job_ids.append(job_id),
            add_interval_job=lambda *, job_id, **kwargs: interval_job_ids.append(job_id),
        ),
    )

    await scheduler.start()

    assert "stats_aggregation" not in cron_job_ids
    assert "audit_cleanup" not in cron_job_ids
    assert "candidate_cleanup" not in cron_job_ids
    assert "db_maintenance" not in cron_job_ids
    assert "stats_hourly_aggregation" not in cron_job_ids
    assert "pool_monitor" not in interval_job_ids
    assert "http_client_idle_cleanup" not in interval_job_ids
    assert "usage_cleanup" not in cron_job_ids
    assert "wallet_daily_usage_aggregation" not in cron_job_ids
    assert scheduler.CHECKIN_JOB_ID not in cron_job_ids
    assert "antigravity_ua_refresh" not in interval_job_ids
    assert "gemini_file_mapping_cleanup" not in interval_job_ids
    assert "pending_cleanup" not in interval_job_ids
    assert "stats_aggregation" not in scheduler._registered_job_ids
    assert "audit_cleanup" not in scheduler._registered_job_ids
    assert "candidate_cleanup" not in scheduler._registered_job_ids
    assert "db_maintenance" not in scheduler._registered_job_ids
    assert "stats_hourly_aggregation" not in scheduler._registered_job_ids
    assert "pool_monitor" not in scheduler._registered_job_ids
    assert "http_client_idle_cleanup" not in scheduler._registered_job_ids
    assert "usage_cleanup" not in scheduler._registered_job_ids
    assert "wallet_daily_usage_aggregation" not in scheduler._registered_job_ids
    assert scheduler.CHECKIN_JOB_ID not in scheduler._registered_job_ids
    assert "antigravity_ua_refresh" not in scheduler._registered_job_ids
    assert "gemini_file_mapping_cleanup" not in scheduler._registered_job_ids
    assert "pending_cleanup" not in scheduler._registered_job_ids


@pytest.mark.asyncio
async def test_maintenance_scheduler_start_registers_python_owned_maintenance_jobs_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(config, "maintenance_startup_tasks_enabled", False)
    monkeypatch.setattr(config, "audit_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "db_maintenance_python_enabled", True)
    monkeypatch.setattr(config, "gemini_file_mapping_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "http_client_idle_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "pending_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "pool_monitor_python_enabled", True)
    monkeypatch.setattr(config, "provider_checkin_python_enabled", True)
    monkeypatch.setattr(config, "request_candidate_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "stats_aggregation_python_enabled", True)
    monkeypatch.setattr(config, "stats_hourly_aggregation_python_enabled", True)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", True)
    monkeypatch.setattr(config, "usage_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "wallet_daily_usage_aggregation_python_enabled", True)

    scheduler = MaintenanceScheduler()
    cron_job_ids: list[str] = []
    interval_job_ids: list[str] = []

    monkeypatch.setattr(scheduler, "_get_checkin_time", lambda: (1, 5))
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "get_scheduler",
        lambda: SimpleNamespace(
            add_cron_job=lambda *, job_id, **kwargs: cron_job_ids.append(job_id),
            add_interval_job=lambda *, job_id, **kwargs: interval_job_ids.append(job_id),
        ),
    )

    await scheduler.start()

    assert "stats_aggregation" in cron_job_ids
    assert "audit_cleanup" in cron_job_ids
    assert "candidate_cleanup" in cron_job_ids
    assert "db_maintenance" in cron_job_ids
    assert "stats_hourly_aggregation" in cron_job_ids
    assert "pool_monitor" in interval_job_ids
    assert "http_client_idle_cleanup" in interval_job_ids
    assert "usage_cleanup" in cron_job_ids
    assert "wallet_daily_usage_aggregation" in cron_job_ids
    assert scheduler.CHECKIN_JOB_ID in cron_job_ids
    assert "antigravity_ua_refresh" in interval_job_ids
    assert "gemini_file_mapping_cleanup" in interval_job_ids
    assert "pending_cleanup" in interval_job_ids
    assert "stats_aggregation" in scheduler._registered_job_ids
    assert "audit_cleanup" in scheduler._registered_job_ids
    assert "candidate_cleanup" in scheduler._registered_job_ids
    assert "db_maintenance" in scheduler._registered_job_ids
    assert "stats_hourly_aggregation" in scheduler._registered_job_ids
    assert "pool_monitor" in scheduler._registered_job_ids
    assert "http_client_idle_cleanup" in scheduler._registered_job_ids
    assert "usage_cleanup" in scheduler._registered_job_ids
    assert "wallet_daily_usage_aggregation" in scheduler._registered_job_ids
    assert scheduler.CHECKIN_JOB_ID in scheduler._registered_job_ids
    assert "antigravity_ua_refresh" in scheduler._registered_job_ids
    assert "gemini_file_mapping_cleanup" in scheduler._registered_job_ids
    assert "pending_cleanup" in scheduler._registered_job_ids


@pytest.mark.asyncio
async def test_maintenance_scheduler_startup_tasks_skip_python_pending_cleanup_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()
    refresh_user_agent = AsyncMock()
    perform_pending_cleanup = AsyncMock()

    monkeypatch.setattr(config, "pending_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", False)
    monkeypatch.setattr(maintenance_scheduler_module.asyncio, "sleep", AsyncMock())
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.client.refresh_user_agent",
        refresh_user_agent,
    )
    monkeypatch.setattr(scheduler, "_perform_pending_cleanup", perform_pending_cleanup)

    await scheduler._run_startup_tasks()

    refresh_user_agent.assert_not_awaited()
    perform_pending_cleanup.assert_not_awaited()


@pytest.mark.asyncio
async def test_maintenance_scheduler_startup_tasks_run_python_pending_cleanup_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()
    refresh_user_agent = AsyncMock()
    perform_pending_cleanup = AsyncMock()

    monkeypatch.setattr(config, "pending_cleanup_python_enabled", True)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", True)
    monkeypatch.setattr(maintenance_scheduler_module.asyncio, "sleep", AsyncMock())
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.client.refresh_user_agent",
        refresh_user_agent,
    )
    monkeypatch.setattr(scheduler, "_perform_pending_cleanup", perform_pending_cleanup)

    await scheduler._run_startup_tasks()

    refresh_user_agent.assert_awaited_once()
    perform_pending_cleanup.assert_awaited_once()


@pytest.mark.asyncio
async def test_start_background_services_skips_python_maintenance_scheduler_when_unused(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, _task_poller, _model_fetch_scheduler = (
        _patch_background_services_dependencies(monkeypatch)
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)
    monkeypatch.setattr(config, "model_fetch_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "audit_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "antigravity_ua_refresh_python_enabled", False)
    monkeypatch.setattr(config, "db_maintenance_python_enabled", False)
    monkeypatch.setattr(config, "gemini_file_mapping_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "http_client_idle_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "pending_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "pool_monitor_python_enabled", False)
    monkeypatch.setattr(config, "provider_checkin_python_enabled", False)
    monkeypatch.setattr(config, "request_candidate_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "stats_aggregation_python_enabled", False)
    monkeypatch.setattr(config, "stats_hourly_aggregation_python_enabled", False)
    monkeypatch.setattr(config, "usage_cleanup_python_enabled", False)
    monkeypatch.setattr(config, "wallet_daily_usage_aggregation_python_enabled", False)

    await main_module._start_background_services(state)

    assert "maintenance_scheduler" not in coordinator.acquire_calls
    assert state.maintenance_scheduler is None


@pytest.mark.asyncio
async def test_start_background_services_starts_python_maintenance_scheduler_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, _task_poller, _model_fetch_scheduler = (
        _patch_background_services_dependencies(
            monkeypatch,
            acquire_results={"maintenance_scheduler": True},
        )
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)
    monkeypatch.setattr(config, "model_fetch_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "http_client_idle_cleanup_python_enabled", True)

    await main_module._start_background_services(state)

    assert "maintenance_scheduler" in coordinator.acquire_calls
    assert "maintenance_scheduler" in coordinator.registered_callbacks
    assert state.maintenance_scheduler is not None


def _patch_core_infrastructure_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(config, "log_startup_warnings", lambda: None)
    monkeypatch.setattr(main_module, "init_db", lambda: None)
    monkeypatch.setattr(main_module, "initialize_providers", AsyncMock())
    monkeypatch.setattr(
        "src.clients.redis_client.get_redis_client",
        AsyncMock(return_value=None),
    )
    monkeypatch.setattr(
        "src.services.rate_limit.concurrency_manager.get_concurrency_manager",
        AsyncMock(return_value=None),
    )
    monkeypatch.setattr(
        "src.services.rate_limit.user_rpm_limiter.get_user_rpm_limiter",
        AsyncMock(return_value=None),
    )
    monkeypatch.setattr("src.core.batch_committer.init_batch_committer", AsyncMock())
    monkeypatch.setattr(
        "src.services.provider_keys.codex_quota_sync_dispatcher.init_codex_quota_sync_dispatcher",
        AsyncMock(),
    )


def _patch_python_host_shutdown_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(
        "src.services.provider_keys.codex_quota_sync_dispatcher.shutdown_codex_quota_sync_dispatcher",
        AsyncMock(),
    )
    monkeypatch.setattr("src.core.batch_committer.shutdown_batch_committer", AsyncMock())
    monkeypatch.setattr("src.clients.redis_client.close_redis_client", AsyncMock())
    monkeypatch.setattr(main_module, "close_http_clients", AsyncMock())


class _FakeStartupTaskCoordinator:
    def __init__(self, acquire_results: dict[str, bool] | None = None) -> None:
        self.acquire_results = acquire_results or {}
        self.acquire_calls: list[str] = []
        self.registered_callbacks: list[str] = []
        self.released: list[str] = []

    async def acquire(self, name: str) -> bool:
        self.acquire_calls.append(name)
        return self.acquire_results.get(name, False)

    async def release(self, name: str) -> None:
        self.released.append(name)

    def register_lock_lost_callback(self, name: str, _callback: Any) -> None:
        self.registered_callbacks.append(name)


def _patch_background_services_dependencies(
    monkeypatch: pytest.MonkeyPatch,
    *,
    acquire_results: dict[str, bool] | None = None,
) -> tuple[_FakeStartupTaskCoordinator, Any, Any, Any]:
    coordinator = _FakeStartupTaskCoordinator(acquire_results)
    quota_scheduler = SimpleNamespace(start=AsyncMock(), stop=AsyncMock())
    maintenance_scheduler = SimpleNamespace(start=AsyncMock(), stop=AsyncMock())
    model_fetch_scheduler = SimpleNamespace(start=AsyncMock(), stop=AsyncMock())
    pool_quota_probe_scheduler = SimpleNamespace(start=AsyncMock(), stop=AsyncMock())
    task_poller = SimpleNamespace(start=AsyncMock(), stop=AsyncMock())
    task_scheduler = SimpleNamespace(start=MagicMock())

    monkeypatch.setattr(
        "src.utils.task_coordinator.StartupTaskCoordinator",
        lambda _redis: coordinator,
    )
    monkeypatch.setattr(
        "src.services.usage.quota_scheduler.get_quota_scheduler",
        lambda: quota_scheduler,
    )
    monkeypatch.setattr(
        "src.services.system.maintenance_scheduler.get_maintenance_scheduler",
        lambda: maintenance_scheduler,
    )
    monkeypatch.setattr(
        "src.services.model.fetch_scheduler.get_model_fetch_scheduler",
        lambda: model_fetch_scheduler,
    )
    monkeypatch.setattr(
        "src.services.provider_keys.pool_quota_probe_scheduler.get_pool_quota_probe_scheduler",
        lambda: pool_quota_probe_scheduler,
    )
    monkeypatch.setattr(
        "src.services.task.polling.task_poller.get_task_poller",
        lambda: task_poller,
    )
    monkeypatch.setattr(
        "src.services.system.scheduler.get_scheduler",
        lambda: task_scheduler,
    )

    return coordinator, quota_scheduler, task_poller, model_fetch_scheduler


@pytest.mark.asyncio
async def test_initialize_core_infrastructure_skips_python_usage_consumer_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _patch_core_infrastructure_dependencies(monkeypatch)
    started = AsyncMock()

    monkeypatch.setattr(config, "require_redis", False)
    monkeypatch.setattr(config, "usage_queue_enabled", True)
    monkeypatch.setattr(config, "usage_queue_python_consumer_enabled", False)
    monkeypatch.setattr("src.services.usage.consumer_streams.start_usage_queue_consumer", started)

    await main_module._initialize_core_infrastructure(main_module.LifecycleState())

    started.assert_not_awaited()


@pytest.mark.asyncio
async def test_initialize_core_infrastructure_starts_python_usage_consumer_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _patch_core_infrastructure_dependencies(monkeypatch)
    started = AsyncMock()

    monkeypatch.setattr(config, "require_redis", False)
    monkeypatch.setattr(config, "usage_queue_enabled", True)
    monkeypatch.setattr(config, "usage_queue_python_consumer_enabled", True)
    monkeypatch.setattr("src.services.usage.consumer_streams.start_usage_queue_consumer", started)

    await main_module._initialize_core_infrastructure(main_module.LifecycleState())

    started.assert_awaited_once()


@pytest.mark.asyncio
async def test_run_python_host_shutdown_skips_python_usage_consumer_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _patch_python_host_shutdown_dependencies(monkeypatch)
    stopped = AsyncMock()

    monkeypatch.setattr(config, "usage_queue_enabled", True)
    monkeypatch.setattr(config, "usage_queue_python_consumer_enabled", False)
    monkeypatch.setattr("src.services.usage.consumer_streams.stop_usage_queue_consumer", stopped)

    await main_module._run_python_host_shutdown(main_module.LifecycleState())

    stopped.assert_not_awaited()


@pytest.mark.asyncio
async def test_run_python_host_shutdown_stops_python_usage_consumer_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    _patch_python_host_shutdown_dependencies(monkeypatch)
    stopped = AsyncMock()

    monkeypatch.setattr(config, "usage_queue_enabled", True)
    monkeypatch.setattr(config, "usage_queue_python_consumer_enabled", True)
    monkeypatch.setattr("src.services.usage.consumer_streams.stop_usage_queue_consumer", stopped)

    await main_module._run_python_host_shutdown(main_module.LifecycleState())

    stopped.assert_awaited_once()


@pytest.mark.asyncio
async def test_start_background_services_skips_python_video_poller_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, task_poller, _model_fetch_scheduler = _patch_background_services_dependencies(
        monkeypatch
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)

    await main_module._start_background_services(state)

    assert "task_poller:video" not in coordinator.acquire_calls
    task_poller.start.assert_not_awaited()
    assert state.task_poller is None


@pytest.mark.asyncio
async def test_start_background_services_starts_python_video_poller_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, task_poller, _model_fetch_scheduler = _patch_background_services_dependencies(
        monkeypatch,
        acquire_results={"task_poller:video": True},
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "video_task_python_poller_enabled", True)

    await main_module._start_background_services(state)

    assert "task_poller:video" in coordinator.acquire_calls
    assert "task_poller:video" in coordinator.registered_callbacks
    task_poller.start.assert_awaited_once()
    assert state.task_poller is task_poller


@pytest.mark.asyncio
async def test_start_background_services_skips_python_quota_scheduler_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, quota_scheduler, _task_poller, _model_fetch_scheduler = _patch_background_services_dependencies(
        monkeypatch
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)

    await main_module._start_background_services(state)

    assert "quota_scheduler" not in coordinator.acquire_calls
    quota_scheduler.start.assert_not_awaited()
    assert state.quota_scheduler is None


@pytest.mark.asyncio
async def test_start_background_services_starts_python_quota_scheduler_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, quota_scheduler, _task_poller, _model_fetch_scheduler = _patch_background_services_dependencies(
        monkeypatch,
        acquire_results={"quota_scheduler": True},
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", True)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)

    await main_module._start_background_services(state)

    assert "quota_scheduler" in coordinator.acquire_calls
    assert "quota_scheduler" in coordinator.registered_callbacks
    quota_scheduler.start.assert_awaited_once()
    assert state.quota_scheduler is quota_scheduler


@pytest.mark.asyncio
async def test_start_background_services_skips_python_model_fetch_scheduler_when_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, _task_poller, model_fetch_scheduler = (
        _patch_background_services_dependencies(monkeypatch)
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "model_fetch_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)

    await main_module._start_background_services(state)

    assert "model_fetch_scheduler" not in coordinator.acquire_calls
    model_fetch_scheduler.start.assert_not_awaited()
    assert state.model_fetch_scheduler is None


@pytest.mark.asyncio
async def test_start_background_services_starts_python_model_fetch_scheduler_when_enabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    coordinator, _quota_scheduler, _task_poller, model_fetch_scheduler = (
        _patch_background_services_dependencies(
            monkeypatch,
            acquire_results={"model_fetch_scheduler": True},
        )
    )
    state = main_module.LifecycleState()

    monkeypatch.setattr(config, "quota_scheduler_python_enabled", False)
    monkeypatch.setattr(config, "model_fetch_scheduler_python_enabled", True)
    monkeypatch.setattr(config, "video_task_python_poller_enabled", False)

    await main_module._start_background_services(state)

    assert "model_fetch_scheduler" in coordinator.acquire_calls
    assert "model_fetch_scheduler" in coordinator.registered_callbacks
    model_fetch_scheduler.start.assert_awaited_once()
    assert state.model_fetch_scheduler is model_fetch_scheduler


@pytest.mark.asyncio
async def test_stop_service_on_lock_lost_keeps_state_when_stop_fails() -> None:
    state = main_module.LifecycleState()
    service = SimpleNamespace()
    state.quota_scheduler = cast(Any, service)

    async def fail_stop() -> None:
        raise RuntimeError("boom")

    await main_module._stop_service_on_lock_lost(
        state,
        lock_name="quota_scheduler",
        service_name="月卡额度重置调度器",
        state_attr="quota_scheduler",
        stop=fail_stop,
    )

    assert state.quota_scheduler is service


@pytest.mark.asyncio
async def test_stop_service_on_lock_lost_clears_state_after_success() -> None:
    state = main_module.LifecycleState()
    service = SimpleNamespace()
    state.quota_scheduler = cast(Any, service)
    stopped = False

    async def stop() -> None:
        nonlocal stopped
        stopped = True

    await main_module._stop_service_on_lock_lost(
        state,
        lock_name="quota_scheduler",
        service_name="月卡额度重置调度器",
        state_attr="quota_scheduler",
        stop=stop,
    )

    assert stopped is True
    assert state.quota_scheduler is None


def test_http_client_idle_cleanup_interval_env_invalid(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("HTTP_CLIENT_IDLE_CLEANUP_INTERVAL_MINUTES", "bad")
    assert MaintenanceScheduler._get_http_client_idle_cleanup_interval_minutes() == 5


@pytest.mark.asyncio
async def test_candidate_cleanup_uses_dedicated_retention_and_batch_settings(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()

    class _FakeLoop:
        async def run_in_executor(self, _executor, func):  # type: ignore[no-untyped-def]
            return func()

    class _ConfigSession:
        def close(self) -> None:
            return None

    class _BatchSession:
        def __init__(self, ids: list[str]) -> None:
            self.ids = ids
            self.closed = False
            self.committed = False
            self.query_obj = MagicMock()
            filtered = self.query_obj.filter.return_value
            filtered.order_by.return_value.limit.return_value.all.return_value = [
                SimpleNamespace(id=value) for value in ids
            ]

        def query(self, _model):  # type: ignore[no-untyped-def]
            return self.query_obj

        def execute(self, _statement):  # type: ignore[no-untyped-def]
            return SimpleNamespace(rowcount=len(self.ids))

        def commit(self) -> None:
            self.committed = True

        def rollback(self) -> None:
            raise AssertionError("rollback should not be called")

        def close(self) -> None:
            self.closed = True

    config_session = _ConfigSession()
    batch_one = _BatchSession(["candidate-1", "candidate-2"])
    batch_two = _BatchSession([])
    sessions = iter([config_session, batch_one, batch_two])

    def fake_create_session():  # type: ignore[no-untyped-def]
        return next(sessions)

    config_values = {
        "enable_auto_cleanup": True,
        "request_candidates_retention_days": 21,
        "request_candidates_cleanup_batch_size": 2,
    }

    monkeypatch.setattr(maintenance_scheduler_module, "create_session", fake_create_session)
    monkeypatch.setattr(
        maintenance_scheduler_module.SystemConfigService,
        "get_config",
        lambda _db, key, default=None: config_values.get(key, default),
    )
    monkeypatch.setattr(
        maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
    )

    await scheduler._perform_candidate_cleanup()

    batch_one.query_obj.filter.return_value.order_by.return_value.limit.return_value.all.assert_called_once()
    batch_one.query_obj.filter.return_value.order_by.return_value.limit.assert_called_once_with(2)
    assert batch_one.committed is True
    assert batch_one.closed is True
    assert batch_two.closed is True


def test_cleanup_body_fields_batches_records_with_single_commit(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()

    class _BatchSession:
        def __init__(self, records: list[SimpleNamespace]) -> None:
            self.records = records
            self.closed = False
            self.committed = False
            self.executed = 0
            self.query_obj = MagicMock()
            filtered = self.query_obj.filter.return_value
            filtered.filter.return_value.order_by.return_value.limit.return_value.all.return_value = (
                records
            )

        def query(self, *args):  # type: ignore[no-untyped-def]
            self.query_args = args
            return self.query_obj

        def execute(self, _statement):  # type: ignore[no-untyped-def]
            self.executed += 1
            return SimpleNamespace(rowcount=1)

        def commit(self) -> None:
            self.committed = True

        def rollback(self) -> None:
            raise AssertionError("rollback should not be called")

        def close(self) -> None:
            self.closed = True

    batch_one = _BatchSession(
        [
            SimpleNamespace(
                id="usage-1",
                request_body={"hello": "world"},
                response_body=None,
                provider_request_body=None,
                client_response_body=None,
            ),
            SimpleNamespace(
                id="usage-2",
                request_body=None,
                response_body={"ok": True},
                provider_request_body=None,
                client_response_body=None,
            ),
        ]
    )
    batch_two = _BatchSession([])
    sessions = iter([batch_one, batch_two])

    monkeypatch.setattr(
        maintenance_scheduler_module,
        "create_session",
        lambda: next(sessions),
    )
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "compress_json",
        lambda payload: f"compressed:{payload}".encode(),
    )

    compressed = scheduler._cleanup_body_fields(
        cutoff_time=SimpleNamespace(),  # type: ignore[arg-type]
        batch_size=1000,
    )

    assert compressed == 2
    assert len(batch_one.query_args) == 5
    assert batch_one.executed == 2
    assert batch_one.committed is True
    assert batch_one.closed is True
    assert batch_two.closed is True


def test_cleanup_header_fields_clears_client_response_headers(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()

    class _BatchSession:
        def __init__(self, ids: list[str]) -> None:
            self.ids = ids
            self.closed = False
            self.committed = False
            self.query_obj = MagicMock()
            self.filtered_by_time = MagicMock()
            self.filtered_by_headers = MagicMock()
            self.query_obj.filter.return_value = self.filtered_by_time
            self.filtered_by_time.filter.return_value = self.filtered_by_headers
            self.filtered_by_headers.order_by.return_value.limit.return_value.all.return_value = [
                SimpleNamespace(id=value) for value in ids
            ]
            self.executed_statements: list[str] = []

        def query(self, *args):  # type: ignore[no-untyped-def]
            self.query_args = args
            return self.query_obj

        def execute(self, statement):  # type: ignore[no-untyped-def]
            self.executed_statements.append(str(statement))
            return SimpleNamespace(rowcount=len(self.ids))

        def commit(self) -> None:
            self.committed = True

        def rollback(self) -> None:
            raise AssertionError("rollback should not be called")

        def close(self) -> None:
            self.closed = True

    batch_one = _BatchSession(["usage-1"])
    batch_two = _BatchSession([])
    sessions = iter([batch_one, batch_two])

    monkeypatch.setattr(
        maintenance_scheduler_module,
        "create_session",
        lambda: next(sessions),
    )

    cleaned = scheduler._cleanup_header_fields(
        cutoff_time=SimpleNamespace(),  # type: ignore[arg-type]
        batch_size=1000,
    )

    header_filter = str(batch_one.filtered_by_time.filter.call_args.args[0])

    assert cleaned == 1
    assert batch_one.query_args == (maintenance_scheduler_module.Usage.id,)
    assert "client_response_headers" in header_filter
    assert "client_response_headers" in batch_one.executed_statements[0]
    assert batch_one.committed is True
    assert batch_one.closed is True
    assert batch_two.closed is True


@pytest.mark.asyncio
async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    scheduler = MaintenanceScheduler()
    fixed_now = datetime(2026, 3, 18, 3, 0, 0, tzinfo=timezone.utc)

    class _FakeDateTime(datetime):
        @classmethod
        def now(cls, tz: timezone | None = None) -> datetime:  # type: ignore[override]
            if tz is None:
                return fixed_now.replace(tzinfo=None)
            return fixed_now.astimezone(tz)

    class _FakeLoop:
        async def run_in_executor(self, _executor, func):  # type: ignore[no-untyped-def]
            return func()

    class _ConfigSession:
        def close(self) -> None:
            return None

    calls: list[tuple[str, datetime, int, datetime | None]] = []

    def _record(name: str, count: int) -> Callable[..., int]:
        def _inner(
            cutoff_time: datetime,
            batch_size: int,
            *,
            newer_than: datetime | None = None,
        ) -> int:
            calls.append((name, cutoff_time, batch_size, newer_than))
            return count

        return _inner

    config_values = {
        "enable_auto_cleanup": True,
        "detail_log_retention_days": 7,
        "compressed_log_retention_days": 30,
        "header_retention_days": 90,
        "log_retention_days": 365,
        "cleanup_batch_size": 123,
        "auto_delete_expired_keys": False,
    }

    def _delete_old_records(cutoff_time: datetime, batch_size: int) -> int:
        calls.append(("delete", cutoff_time, batch_size, None))
        return 5

    monkeypatch.setattr(maintenance_scheduler_module, "datetime", _FakeDateTime)
    monkeypatch.setattr(
        maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
    )
    monkeypatch.setattr(
        maintenance_scheduler_module,
        "create_session",
        lambda: _ConfigSession(),
    )
    monkeypatch.setattr(
        maintenance_scheduler_module.SystemConfigService,
        "get_config",
        lambda _db, key, default=None: config_values.get(key, default),
    )
    monkeypatch.setattr(
        scheduler,
        "_delete_old_records",
        _delete_old_records,
    )
    monkeypatch.setattr(
        scheduler,
        "_cleanup_header_fields",
        _record("header", 4),
    )
    monkeypatch.setattr(
        scheduler,
        "_cleanup_stale_body_fields",
        _record("body", 3),
    )
    monkeypatch.setattr(
        scheduler,
        "_cleanup_body_fields",
        _record("compress", 2),
    )
    monkeypatch.setattr(
        maintenance_scheduler_module.ApiKeyService,
        "cleanup_expired_keys",
        lambda _db, auto_delete=False: 0,
    )

    await scheduler._perform_cleanup()

    detail_cutoff = fixed_now - timedelta(days=7)
    compressed_cutoff = fixed_now - timedelta(days=30)
    header_cutoff = fixed_now - timedelta(days=90)
    log_cutoff = fixed_now - timedelta(days=365)

    assert calls == [
        ("delete", log_cutoff, 123, None),
        ("header", header_cutoff, 123, log_cutoff),
        ("body", compressed_cutoff, 123, log_cutoff),
        ("compress", detail_cutoff, 123, compressed_cutoff),
    ]