Jonathan
105 lines
4.7 KiB
Python
105 lines
4.7 KiB
Python
from lyra_worker import main
|
|
from lyra_worker.adapters.fakes import FakeMbBrowser, FakeSimilaritySource
|
|
from lyra_worker.discovery import DiscoveryConfig, DiscoveryResult
|
|
from lyra_worker.main import maybe_run_discovery
|
|
from lyra_worker.similarity.base import SimilarArtist
|
|
from tests.conftest import insert_watched_artist
|
|
|
|
|
|
def _set(conn, key, value):
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'INSERT INTO "Config"(key, value, secret, "updatedAt") VALUES (%s, %s, false, now()) '
|
|
'ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, "updatedAt" = now()',
|
|
(key, value),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _get(conn, key):
|
|
with conn.cursor() as cur:
|
|
cur.execute('SELECT value FROM "Config" WHERE key = %s', (key,))
|
|
row = cur.fetchone()
|
|
return row[0] if row else None
|
|
|
|
|
|
def test_maybe_run_discovery_drains_across_chunks_and_finishes(conn):
|
|
# The sweep must advance one chunk (discover.chunkSize seeds) per call, clear
|
|
# discover.requested exactly once at start, and finally drain on the 0-seed chunk:
|
|
# writing discover.result and clearing discover.inProgress.
|
|
for i in range(3):
|
|
insert_watched_artist(conn, mbid=f"s{i}", name=f"Seed{i}")
|
|
_set(conn, "discover.requested", "true")
|
|
_set(conn, "discover.chunkSize", "2")
|
|
src = FakeSimilaritySource(
|
|
similar={f"s{i}": [SimilarArtist(f"c{i}", f"C{i}", 0.9)] for i in range(3)}
|
|
)
|
|
cfg = {"discover.enabled": "false", "discover.requested": "true", "discover.chunkSize": "2"}
|
|
tick = 0.0
|
|
# iteration 1: starts sweep, processes 2 seeds, still in progress
|
|
tick = maybe_run_discovery(
|
|
conn, [src], FakeMbBrowser(), cfg, DiscoveryConfig.from_config(cfg), 1.0, tick
|
|
)
|
|
assert _get(conn, "discover.inProgress") == "true"
|
|
assert _get(conn, "discover.requested") == "false" # cleared at sweep start
|
|
# refresh the snapshot the way run_forever's loop does (re-reads config each iteration)
|
|
cfg2 = {**cfg, "discover.requested": "false", "discover.inProgress": "true",
|
|
"discover.progress": _get(conn, "discover.progress")}
|
|
# iteration 2: processes the last seed, still in progress
|
|
maybe_run_discovery(
|
|
conn, [src], FakeMbBrowser(), cfg2, DiscoveryConfig.from_config(cfg2), 2.0, tick
|
|
)
|
|
assert _get(conn, "discover.inProgress") == "true"
|
|
cfg3 = {**cfg2, "discover.progress": _get(conn, "discover.progress")}
|
|
# iteration 3: 0 seeds left -> drains, writes result, clears inProgress
|
|
maybe_run_discovery(
|
|
conn, [src], FakeMbBrowser(), cfg3, DiscoveryConfig.from_config(cfg3), 3.0, tick
|
|
)
|
|
assert _get(conn, "discover.inProgress") == "false"
|
|
assert _get(conn, "discover.progress") == ""
|
|
assert "artists" in (_get(conn, "discover.result") or "")
|
|
|
|
|
|
def test_maybe_run_discovery_runs_once_when_due_and_requested(conn, monkeypatch):
|
|
# A scheduled tick that is due AND a pending one-shot request in the same loop
|
|
# iteration must trigger discovery exactly once, not twice.
|
|
calls = []
|
|
monkeypatch.setattr(main, "_run_discovery",
|
|
lambda *a, **k: (calls.append(1), DiscoveryResult(0, 0, seeds=0))[1])
|
|
config = {"discover.enabled": "true", "discover.requested": "true"}
|
|
dcfg = DiscoveryConfig.from_config(config)
|
|
new_tick = maybe_run_discovery(
|
|
conn, [], FakeMbBrowser(), config, dcfg, now=10_000.0, last_discover_tick=0.0
|
|
)
|
|
assert calls == [1] # exactly one run
|
|
assert new_tick == 10_000.0 # schedule tick advanced
|
|
|
|
|
|
def test_maybe_run_discovery_requested_when_disabled_does_not_advance_tick(conn, monkeypatch):
|
|
# A one-shot request runs even when discovery is disabled, but must not advance the
|
|
# schedule tick (there is no active schedule).
|
|
calls = []
|
|
monkeypatch.setattr(main, "_run_discovery",
|
|
lambda *a, **k: (calls.append(1), DiscoveryResult(0, 0, seeds=0))[1])
|
|
config = {"discover.enabled": "false", "discover.requested": "true"}
|
|
dcfg = DiscoveryConfig.from_config(config)
|
|
new_tick = maybe_run_discovery(
|
|
conn, [], FakeMbBrowser(), config, dcfg, now=10_000.0, last_discover_tick=42.0
|
|
)
|
|
assert calls == [1]
|
|
assert new_tick == 42.0 # unchanged
|
|
|
|
|
|
def test_maybe_run_discovery_idle_does_nothing(conn, monkeypatch):
|
|
calls = []
|
|
monkeypatch.setattr(main, "_run_discovery",
|
|
lambda *a, **k: (calls.append(1), DiscoveryResult(0, 0, seeds=0))[1])
|
|
config = {"discover.enabled": "true", "discover.requested": "false"}
|
|
dcfg = DiscoveryConfig.from_config(config)
|
|
# enabled but not yet due (tick was recent relative to now)
|
|
new_tick = maybe_run_discovery(
|
|
conn, [], FakeMbBrowser(), config, dcfg, now=100.0, last_discover_tick=99.0
|
|
)
|
|
assert calls == []
|
|
assert new_tick == 99.0
|