91c739c2af
Add tests/test_allowlist.py and tests/test_admin_allowlist.py covering config validation, poller filter/prune behavior, admin CRUD, and integration scenarios A/B/C plus NeuralWatt regression. Also set conn.row_factory = sqlite3.Row in poller.main() so the new allowlist SELECT returns Row objects, matching the rest of the codebase.
631 lines
20 KiB
Python
631 lines
20 KiB
Python
"""Tests for the opt-in provider model allowlist.
|
|
|
|
Covers config validation, poller filtering / pruning, and end-to-end
|
|
integration scenarios. All tests are offline: the poller's ``requests.get``
|
|
is mocked.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import sqlite3
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import poller
|
|
|
|
from config import DispatchProvider, RouterConfig, load_config
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
SCHEMA_SQL = (ROOT / "config" / "schema.sql").read_text()
|
|
REAL_CFG = load_config(str(ROOT / "config" / "config.yaml"))
|
|
|
|
# Exact 30 model IDs shipped in config/allowlist-openrouter.sql.
|
|
ALLOWLIST_MODEL_IDS = [
|
|
"mistralai/codestral-2508",
|
|
"mistralai/devstral-2512",
|
|
"qwen/qwen3-coder",
|
|
"qwen/qwen3-coder-plus",
|
|
"qwen/qwen3-coder-flash",
|
|
"qwen/qwen3-coder-next",
|
|
"moonshotai/kimi-k2.7-code",
|
|
"z-ai/glm-5.3",
|
|
"moonshotai/kimi-k3",
|
|
"qwen/qwen3-max-thinking",
|
|
"google/gemini-3.1-pro-preview",
|
|
"deepseek/deepseek-v4-flash",
|
|
"deepseek/deepseek-v4-pro",
|
|
"z-ai/glm-5.3-flash",
|
|
"z-ai/glm-5.2",
|
|
"qwen/qwen3-vl-235b-a22b-instruct",
|
|
"qwen/qwen3-vl-30b-a3b-instruct",
|
|
"google/gemini-3.5-flash",
|
|
"google/gemini-3.1-flash-lite",
|
|
"qwen/qwen3.6-35b-a3b",
|
|
"qwen/qwen3-max",
|
|
"mistralai/mistral-medium-3.1",
|
|
"mistralai/mistral-small-3.2-24b-instruct",
|
|
"mistralai/mistral-large-2512",
|
|
"z-ai/glm-5.2:free",
|
|
"nvidia/nemotron-3-super-120b-a12b:free",
|
|
"nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free",
|
|
"meta-llama/llama-4-maverick",
|
|
"meta-llama/llama-4-scout",
|
|
"deepseek/deepseek-v3.2",
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _fake_response(json_data: dict):
|
|
fake = MagicMock()
|
|
fake.json.return_value = json_data
|
|
fake.raise_for_status.return_value = None
|
|
fake.status_code = 200
|
|
return fake
|
|
|
|
|
|
def _catalog_item(model_id: str) -> dict:
|
|
return {
|
|
"id": model_id,
|
|
"metadata": {
|
|
"display_name": model_id.replace("-", " ").title(),
|
|
"description": "test model",
|
|
"pricing": {
|
|
"input_per_million": 0.30,
|
|
"output_per_million": 0.60,
|
|
"cached_input_per_million": 0.20,
|
|
"pricing_tbd": False,
|
|
},
|
|
"capabilities": {
|
|
"tools": True,
|
|
"json_mode": True,
|
|
"vision": False,
|
|
"reasoning": False,
|
|
},
|
|
"limits": {
|
|
"max_context_length": 131072,
|
|
"max_output_tokens": 8192,
|
|
},
|
|
"deprecated": False,
|
|
},
|
|
}
|
|
|
|
|
|
def _openrouter_model(model_id: str) -> dict:
|
|
return {
|
|
"id": model_id,
|
|
"name": model_id.replace("/", " ").title(),
|
|
"canonical_slug": model_id.split(":")[0],
|
|
"pricing": {"prompt": "0.000001", "completion": "0.000002"},
|
|
"top_provider": {"context_length": 131072, "max_completion_tokens": 8192},
|
|
"architecture": {"input_modalities": ["text"]},
|
|
"supported_parameters": [],
|
|
"reasoning": {},
|
|
}
|
|
|
|
|
|
def _openrouter_catalog(model_ids: list[str]) -> list[dict]:
|
|
return [_openrouter_model(mid) for mid in model_ids]
|
|
|
|
|
|
def _seed_allowlist(conn: sqlite3.Connection, provider: str = "openrouter") -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
for mid in ALLOWLIST_MODEL_IDS:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO provider_model_allowlist "
|
|
"(provider, model_id, added_at, note) VALUES (?, ?, ?, ?)",
|
|
(provider, mid, now, "seed"),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _seed_existing_models(
|
|
conn: sqlite3.Connection,
|
|
model_ids: list[str],
|
|
provider: str = "openrouter",
|
|
deprecated: bool = False,
|
|
) -> None:
|
|
ts = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat()
|
|
availability = "deprecated" if deprecated else "active"
|
|
for mid in model_ids:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO models (
|
|
model_id, provider, base_model_id, display_name,
|
|
cost_per_1m_prompt, cost_per_1m_completion, cost_per_1m_prompt_cached,
|
|
context_window, effective_context_window, max_output_tokens,
|
|
supports_tools, supports_json_mode, supports_vision, supports_reasoning,
|
|
reasoning_default_enabled, latency_class, reasoning_mode,
|
|
context_variant, access_level, pricing_tbd, deprecated,
|
|
availability, last_updated
|
|
) VALUES (?, ?, ?, ?,
|
|
0.30, 0.60, NULL,
|
|
131072, 192500, 8192,
|
|
1, 1, 0, 0,
|
|
0, 'standard', 'none', 'standard', 'public', 0, ?,
|
|
?, ?)
|
|
""",
|
|
(mid, provider, mid, mid, int(deprecated), availability, ts),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _allowlist_cfg(tmp_path: Path, *, openrouter: bool = True) -> RouterConfig:
|
|
"""Return a config with the database pointing at a fresh test DB.
|
|
|
|
OpenRouter keeps ``require_allowlist: true`` from config.yaml. NeuralWatt
|
|
keeps ``require_allowlist: false``.
|
|
"""
|
|
db_path = str(tmp_path / "test.db")
|
|
providers = {}
|
|
for name, prov_cfg in REAL_CFG.dispatch_providers.items():
|
|
if name == "openrouter" and not openrouter:
|
|
providers[name] = prov_cfg.model_copy(update={"enabled": False})
|
|
else:
|
|
providers[name] = prov_cfg
|
|
return REAL_CFG.model_copy(
|
|
update={
|
|
"database": REAL_CFG.database.model_copy(update={"path": db_path}),
|
|
"freshness": REAL_CFG.freshness.model_copy(update={"stale_after_days": 300}),
|
|
"dispatch_providers": providers,
|
|
}
|
|
)
|
|
|
|
|
|
def _openrouter_only_cfg(tmp_path: Path) -> RouterConfig:
|
|
"""Return a config that enables only the openrouter provider."""
|
|
db_path = str(tmp_path / "test.db")
|
|
or_cfg = REAL_CFG.dispatch_providers["openrouter"]
|
|
return REAL_CFG.model_copy(
|
|
update={
|
|
"database": REAL_CFG.database.model_copy(update={"path": db_path}),
|
|
"freshness": REAL_CFG.freshness.model_copy(update={"stale_after_days": 300}),
|
|
"dispatch_providers": {"openrouter": or_cfg},
|
|
}
|
|
)
|
|
|
|
|
|
def _run_poller(cfg: RouterConfig, monkeypatch, catalog_by_url: dict[str, dict]):
|
|
"""Run poller.main() with a URL-keyed mock and return a connect function."""
|
|
db_path = cfg.database.path
|
|
|
|
def _connect():
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
c = _connect()
|
|
c.executescript(SCHEMA_SQL)
|
|
c.close()
|
|
|
|
monkeypatch.setattr(poller, "load_config", lambda p: cfg)
|
|
|
|
def _get(url, timeout, **kwargs):
|
|
for key, payload in catalog_by_url.items():
|
|
if key in url:
|
|
return _fake_response(payload)
|
|
raise RuntimeError(f"unexpected URL in test: {url}")
|
|
|
|
monkeypatch.setattr(poller.requests, "get", _get)
|
|
return _connect
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_dispatch_provider_defaults_require_allowlist_false():
|
|
"""``require_allowlist`` defaults to False when not supplied."""
|
|
provider = DispatchProvider(
|
|
base_url="https://example.com/v1",
|
|
api_key_env="EXAMPLE_API_KEY",
|
|
enabled=True,
|
|
)
|
|
assert provider.require_allowlist is False
|
|
|
|
|
|
def test_dispatch_provider_accepts_require_allowlist_true(tmp_path):
|
|
"""A provider block with ``require_allowlist: true`` loads cleanly."""
|
|
cfg_path = tmp_path / "config.yaml"
|
|
providers = {
|
|
name: prov.model_copy(update={"require_allowlist": True})
|
|
for name, prov in REAL_CFG.dispatch_providers.items()
|
|
}
|
|
cfg_data = REAL_CFG.model_copy(
|
|
update={"dispatch_providers": providers}
|
|
).model_dump(mode="json")
|
|
cfg_path.write_text(json.dumps(cfg_data))
|
|
cfg = load_config(str(cfg_path))
|
|
for prov_cfg in cfg.dispatch_providers.values():
|
|
assert prov_cfg.require_allowlist is True
|
|
|
|
|
|
def test_dispatch_provider_unknown_key_rejected():
|
|
"""``StrictModel.extra="forbid"`` rejects an unknown allowlist-ish key."""
|
|
with pytest.raises(Exception):
|
|
DispatchProvider(
|
|
base_url="https://example.com/v1",
|
|
api_key_env="KEY",
|
|
allowed_models=["foo"],
|
|
)
|
|
|
|
|
|
def test_real_config_openrouter_requires_allowlist():
|
|
"""The shipped config keeps the OpenRouter provider in allowlist mode."""
|
|
assert REAL_CFG.dispatch_providers["openrouter"].require_allowlist is True
|
|
|
|
|
|
def test_real_config_neuralwatt_no_allowlist():
|
|
"""The NeuralWatt provider in the shipped config is not allowlisted."""
|
|
nw = REAL_CFG.dispatch_providers["neuralwatt"]
|
|
assert nw.require_allowlist is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Poller tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_filter_before_upsert(tmp_path, monkeypatch):
|
|
"""Only allowlisted rows are upserted; pruned catalog rows never hit the DB."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
extras = [f"extra/model-{i}" for i in range(5)]
|
|
catalog = {"data": _openrouter_catalog(ALLOWLIST_MODEL_IDS + extras)}
|
|
|
|
connect = _run_poller(cfg, monkeypatch, {"openrouter.ai": catalog})
|
|
c = connect()
|
|
_seed_allowlist(c)
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active_ids = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='active'"
|
|
)
|
|
}
|
|
all_ids = {
|
|
row["model_id"]
|
|
for row in c.execute("SELECT model_id FROM models WHERE provider='openrouter'")
|
|
}
|
|
assert active_ids == set(ALLOWLIST_MODEL_IDS)
|
|
assert all_ids == set(ALLOWLIST_MODEL_IDS)
|
|
for extra in extras:
|
|
assert extra not in all_ids
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_prune_non_allowlisted_existing_rows(tmp_path, monkeypatch):
|
|
"""Pre-existing active rows outside the allowlist become deprecated."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
existing = ALLOWLIST_MODEL_IDS[:5] + [f"legacy/model-{i}" for i in range(10)]
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _openrouter_catalog(ALLOWLIST_MODEL_IDS)}},
|
|
)
|
|
c = connect()
|
|
_seed_allowlist(c)
|
|
_seed_existing_models(c, existing, provider="openrouter")
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='active'"
|
|
)
|
|
}
|
|
deprecated = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
)
|
|
}
|
|
assert active == set(ALLOWLIST_MODEL_IDS)
|
|
assert all(mid in deprecated for mid in existing[5:])
|
|
assert all(mid in active for mid in ALLOWLIST_MODEL_IDS[:5])
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_empty_allowlist_deprecates_all_and_skips_fetch(tmp_path, monkeypatch):
|
|
"""With an empty allowlist all active rows are deprecated and no catalog rows load."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
existing = [f"legacy/model-{i}" for i in range(5)]
|
|
|
|
stderr = io.StringIO()
|
|
monkeypatch.setattr("sys.stderr", stderr)
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _openrouter_catalog(ALLOWLIST_MODEL_IDS)}},
|
|
)
|
|
c = connect()
|
|
_seed_existing_models(c, existing, provider="openrouter")
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
stderr_value = stderr.getvalue()
|
|
assert "allowlist is empty" in stderr_value.lower()
|
|
|
|
c = connect()
|
|
try:
|
|
active_count = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='active'"
|
|
).fetchone()[0]
|
|
deprecated = [
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
)
|
|
]
|
|
assert active_count == 0
|
|
assert set(deprecated) == set(existing)
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_require_allowlist_false_skips_logic(tmp_path, monkeypatch):
|
|
"""NeuralWatt (require_allowlist=False) upserts every fetched row untouched."""
|
|
cfg = _allowlist_cfg(tmp_path, openrouter=False)
|
|
model_ids = ["nw-model-a", "nw-model-b", "nw-model-c"]
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"api.neuralwatt.com": {"data": [_catalog_item(mid) for mid in model_ids]}},
|
|
)
|
|
c = connect()
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='neuralwatt' AND availability='active'"
|
|
)
|
|
}
|
|
deprecated = [
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='neuralwatt' AND deprecated=1"
|
|
)
|
|
]
|
|
assert active == set(model_ids)
|
|
assert deprecated == []
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_add_remove_mid_cycle(tmp_path, monkeypatch):
|
|
"""Changing the allowlist between polls is reflected in the next run."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
first_allowed = ALLOWLIST_MODEL_IDS[:10]
|
|
second_allowed = ALLOWLIST_MODEL_IDS[5:15]
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _openrouter_catalog(ALLOWLIST_MODEL_IDS[:15])}},
|
|
)
|
|
c = connect()
|
|
_seed_allowlist(c)
|
|
c.execute(
|
|
"DELETE FROM provider_model_allowlist WHERE provider='openrouter' AND model_id NOT IN (%s)"
|
|
% ",".join("?" * len(first_allowed)),
|
|
first_allowed,
|
|
)
|
|
c.commit()
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active_first = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='active'"
|
|
)
|
|
}
|
|
assert active_first == set(first_allowed)
|
|
finally:
|
|
c.close()
|
|
|
|
c = connect()
|
|
c.execute("DELETE FROM provider_model_allowlist WHERE provider='openrouter'")
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
for mid in second_allowed:
|
|
c.execute(
|
|
"INSERT INTO provider_model_allowlist (provider, model_id, added_at, note) VALUES (?, ?, ?, ?)",
|
|
("openrouter", mid, now, "mid-cycle"),
|
|
)
|
|
c.commit()
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active_second = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='active'"
|
|
)
|
|
}
|
|
deprecated_now = {
|
|
row["model_id"]
|
|
for row in c.execute(
|
|
"SELECT model_id FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
)
|
|
}
|
|
assert active_second == set(second_allowed)
|
|
assert all(mid in deprecated_now for mid in first_allowed[:5])
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _catalog_425() -> list[dict]:
|
|
"""Return a 425-row OpenRouter-style catalog containing the 30 allowlist IDs."""
|
|
extras = [f"vendor/extra-{i}" for i in range(395)]
|
|
return _openrouter_catalog(ALLOWLIST_MODEL_IDS + extras)
|
|
|
|
|
|
def test_scenario_a_clean_db_allowlist_30_of_425(tmp_path, monkeypatch):
|
|
"""Clean DB + 30 allowlist rows + 425-row catalog -> 30 active, 0 deprecated."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _catalog_425()}},
|
|
)
|
|
c = connect()
|
|
_seed_allowlist(c)
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='active'"
|
|
).fetchone()[0]
|
|
deprecated = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
).fetchone()[0]
|
|
total = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter'"
|
|
).fetchone()[0]
|
|
assert active == 30
|
|
assert deprecated == 0
|
|
assert total == 30
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_scenario_b_preexisting_425_pruned_to_30(tmp_path, monkeypatch):
|
|
"""425 pre-existing rows + 30 allowlist rows -> 30 active + 395 deprecated."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
all_ids = ALLOWLIST_MODEL_IDS + [f"vendor/extra-{i}" for i in range(395)]
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _catalog_425()}},
|
|
)
|
|
c = connect()
|
|
_seed_allowlist(c)
|
|
_seed_existing_models(c, all_ids, provider="openrouter")
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='active'"
|
|
).fetchone()[0]
|
|
deprecated = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
).fetchone()[0]
|
|
total = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter'"
|
|
).fetchone()[0]
|
|
assert active == 30
|
|
assert deprecated == 395
|
|
assert total == 425
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_scenario_c_empty_allowlist_deprecates_all_preexisting(tmp_path, monkeypatch):
|
|
"""Empty allowlist + pre-existing rows -> 0 active, all deprecated, warning logged."""
|
|
cfg = _openrouter_only_cfg(tmp_path)
|
|
existing = ALLOWLIST_MODEL_IDS + [f"vendor/extra-{i}" for i in range(10)]
|
|
|
|
stderr = io.StringIO()
|
|
monkeypatch.setattr("sys.stderr", stderr)
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"openrouter.ai": {"data": _catalog_425()}},
|
|
)
|
|
c = connect()
|
|
_seed_existing_models(c, existing, provider="openrouter")
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
stderr_value = stderr.getvalue()
|
|
assert "allowlist is empty" in stderr_value.lower()
|
|
|
|
c = connect()
|
|
try:
|
|
active = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='active'"
|
|
).fetchone()[0]
|
|
deprecated = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='openrouter' AND availability='deprecated'"
|
|
).fetchone()[0]
|
|
assert active == 0
|
|
assert deprecated == len(existing)
|
|
finally:
|
|
c.close()
|
|
|
|
|
|
def test_neuralwatt_regression_no_deprecation(tmp_path, monkeypatch):
|
|
"""NeuralWatt (allowlist disabled) still marks every fetched row active."""
|
|
cfg = _allowlist_cfg(tmp_path, openrouter=False)
|
|
model_ids = [f"nw-model-{i}" for i in range(20)]
|
|
|
|
connect = _run_poller(
|
|
cfg,
|
|
monkeypatch,
|
|
{"api.neuralwatt.com": {"data": [_catalog_item(mid) for mid in model_ids]}},
|
|
)
|
|
c = connect()
|
|
c.close()
|
|
|
|
poller.main()
|
|
|
|
c = connect()
|
|
try:
|
|
active = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='neuralwatt' AND availability='active'"
|
|
).fetchone()[0]
|
|
deprecated = c.execute(
|
|
"SELECT COUNT(*) FROM models WHERE provider='neuralwatt' AND availability='deprecated'"
|
|
).fetchone()[0]
|
|
assert active == 20
|
|
assert deprecated == 0
|
|
finally:
|
|
c.close()
|