79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
from datetime import UTC, datetime, timedelta
|
|
from unittest.mock import Mock
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from modelforge_api.persistence.models import Base, Model, ModelRevision, ModelUpdateWatch
|
|
from modelforge_api.services.model_update_polling import ModelUpdatePollingService
|
|
from modelforge_api.services.registry import RegistryConflict
|
|
from modelforge_api.settings import Settings
|
|
|
|
|
|
def test_monitor_checks_only_one_due_enabled_unclaimed_watch_and_passes_due_fence() -> None:
|
|
engine = create_engine("sqlite://")
|
|
Base.metadata.create_all(engine)
|
|
settings = Settings(_env_file=None, env="test")
|
|
now = datetime.now(UTC)
|
|
with Session(engine) as session:
|
|
model = Model(key="chosen", display_name="Chosen", upstream_provider="huggingface",
|
|
upstream_source="fixture/chosen")
|
|
session.add(model)
|
|
session.flush()
|
|
revision = ModelRevision(model_id=model.id, upstream_revision="main",
|
|
resolved_commit_sha="a" * 40)
|
|
session.add(revision)
|
|
session.flush()
|
|
watches = [ModelUpdateWatch(
|
|
model_id=model.id, baseline_revision_id=revision.id,
|
|
baseline_commit_sha=revision.resolved_commit_sha, upstream_ref=ref,
|
|
enabled=enabled, last_checked_at=checked, check_until=claimed,
|
|
) for ref, enabled, checked, claimed in (
|
|
("paused", False, None, None),
|
|
("fresh", True, now, None),
|
|
("busy", True, None, now + timedelta(minutes=5)),
|
|
("due", True, now - timedelta(days=2), None),
|
|
("later", True, now - timedelta(hours=25), None),
|
|
)]
|
|
session.add_all(watches)
|
|
session.commit()
|
|
due_id = watches[3].id
|
|
poller = ModelUpdatePollingService(settings, engine)
|
|
service = Mock()
|
|
poller.service = Mock(return_value=service)
|
|
poller.run_once()
|
|
assert service.check.call_count == 1
|
|
assert service.check.call_args.args == (due_id,)
|
|
cutoff = service.check.call_args.kwargs["if_checked_before"]
|
|
assert now - timedelta(days=1) <= cutoff <= datetime.now(UTC) - timedelta(days=1)
|
|
# A concurrent winner is a safe no-op and cannot kill the monitor.
|
|
service.check.side_effect = RegistryConflict("already checked")
|
|
poller.run_once()
|
|
poller.stop()
|
|
poller.run_once()
|
|
assert service.check.call_count == 2
|
|
engine.dispose()
|
|
|
|
|
|
def test_disabled_monitor_never_opens_database_or_constructs_provider() -> None:
|
|
engine = Mock()
|
|
poller = ModelUpdatePollingService(
|
|
Settings(_env_file=None, env="test", model_update_monitor_enabled=False), engine,
|
|
)
|
|
poller.service = Mock()
|
|
poller.run_once()
|
|
assert engine.mock_calls == []
|
|
poller.service.assert_not_called()
|
|
|
|
|
|
def test_unexpected_database_failure_is_redacted_and_does_not_escape(capsys) -> None:
|
|
# Deliberately absent tables in an isolated database; no deployment state is consulted.
|
|
engine = create_engine("sqlite://")
|
|
poller = ModelUpdatePollingService(Settings(_env_file=None, env="test"), engine)
|
|
poller.run_once()
|
|
output = capsys.readouterr().out
|
|
assert "model_update_monitor_check_failed" in output
|
|
assert "OperationalError" in output
|
|
assert "SELECT" not in output
|
|
engine.dispose()
|