muse/tests/test_acquisition.py

"""The acquisition ladder: query variants, candidate fallback, stall handling.

This is the path that had no coverage at all, and it is where every production
failure lived. A fake slskd stands in for the network so the branches that only
happen on a bad day -- a peer that rejects, a transfer slskd forgets, a search
the server filters -- can be exercised deterministically.

Run standalone: python tests/test_acquisition.py
"""
import asyncio, json, os, sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault("MUSE_AUTH", "off")
os.environ.setdefault("MUSE_DB", "dev/acquisition.db")
os.environ.setdefault("MUSE_SEARCH_POLLS", "2")

for f in ("dev/acquisition.db", "dev/acquisition.db-wal", "dev/acquisition.db-shm"):
    if os.path.exists(f):
        os.remove(f)

import app  # noqa: E402

ok = fail = 0


def check(name, cond, detail=""):
    global ok, fail
    if cond:
        ok += 1
        print(f"  PASS  {name}")
    else:
        fail += 1
        print(f"  FAIL  {name}  {detail}")


# Searches poll on a 2s cadence; nothing here is actually waiting on a network.
async def _instant(_):
    return None


app.asyncio.sleep = _instant
app.init_db()          # normally done in lifespan; there is no server here


# ── a slskd that behaves however the test needs ──────────────────────────

class Resp:
    def __init__(self, data=None, status=200):
        self._data, self.status = data, status

    def raise_for_status(self):
        if self.status >= 400:
            raise RuntimeError(f"Server error '{self.status} Internal Server Error'")

    def json(self):
        return self._data


class FakeSlskd:
    """Records what was asked of it and answers from a script."""

    def __init__(self, per_query=None, enqueue_status=None, transfers=None):
        self.per_query = per_query or {}      # query text -> list of responses
        self.enqueue_status = enqueue_status or {}   # username -> http status
        self.transfers = transfers or []
        self.searched = []
        self.enqueued = []
        self.deleted = []
        self._last = None

    def __call__(self, timeout=30.0):
        return self

    async def __aenter__(self):
        return self

    async def __aexit__(self, *a):
        return False

    async def post(self, url, json=None):
        if url == "/api/v0/searches":
            self._last = json["searchText"]
            self.searched.append(self._last)
            return Resp({})
        user = url.rsplit("/", 1)[-1]
        st = self.enqueue_status.get(user, 200)
        self.enqueued.append(user)
        return Resp({}, st)

    async def get(self, url, params=None):
        if url.startswith("/api/v0/searches/"):
            responses = self.per_query.get(self._last, [])
            return Resp({"state": "Completed, ResponseLimitReached",
                         "responses": responses,
                         "fileCount": sum(len(r["files"]) for r in responses)})
        return Resp(self.transfers)

    async def delete(self, url, params=None):
        self.deleted.append(url)
        return Resp({})


def peer(username, filename, *, slot=True, queue=0, size=5_000_000, bitrate=320):
    return {"username": username, "hasFreeUploadSlot": slot,
            "queueLength": queue, "uploadSpeed": 1_000_000,
            "files": [{"filename": filename, "size": size, "bitRate": bitrate}]}


def make_job(job_id, artist="Michael Jackson", title="Billie Jean"):
    with app.db() as c:
        c.execute("INSERT OR REPLACE INTO jobs (id, state, artist, title, album)"
                  " VALUES (?,?,?,?,?)", (job_id, "searching", artist, title, None))


def job_row(job_id):
    with app.db() as c:
        return c.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()


def main() -> int:
    print("\nnormalisation")
    check("accents fold to ASCII", app.norm("Beyoncé") == "beyonce",
          repr(app.norm("Beyoncé")))
    check("umlauts fold", app.norm("Björk") == "bjork", repr(app.norm("Björk")))
    check("slashed o folds", app.norm("margø") == "margo", repr(app.norm("margø")))
    check("apostrophes vanish, not split",
          app.norm("Don't Stop") == "dont stop", repr(app.norm("Don't Stop")))
    check("ampersand still separates", app.norm("R&B") == "r b", repr(app.norm("R&B")))

    print("\nquery ladder")
    qs = app.search_queries("Michael Jackson", "Billie Jean")
    check("falls back to the bare title", "Billie Jean" in qs, qs)
    check("full query is still tried first", qs[0] == "Michael Jackson Billie Jean", qs)
    qs = app.search_queries("Frank Ocean", "Novacane (Explicit Version)")
    check("strips a parenthetical suffix", "Frank Ocean Novacane" in qs, qs)
    qs = app.search_queries("Post Sex Nachos", "Pop Tarts (feat. MARIS)")
    check("strips a featuring credit", "Post Sex Nachos Pop Tarts" in qs, qs)
    qs = app.search_queries("The Weeknd", "House Of Balloons / Glass Table Girls")
    check("splits a slashed title", "House Of Balloons" in qs, qs)
    check("ladder is capped", len(app.search_queries("A", "B")) <= app.MAX_QUERIES)

    print("\nscoring")
    avail_mp3 = app.score({"filename": r"x\Radiohead - Creep.mp3", "size": 1,
                           "bitRate": 320},
                          {"username": "u", "hasFreeUploadSlot": True,
                           "queueLength": 0, "uploadSpeed": 0},
                          "Radiohead", "Creep")
    queued_flac = app.score({"filename": r"x\Radiohead - Creep.flac", "size": 1,
                             "bitRate": 1000},
                            {"username": "v", "hasFreeUploadSlot": False,
                             "queueLength": 400, "uploadSpeed": 0},
                            "Radiohead", "Creep")
    check("an available MP3 beats a FLAC behind 400 people",
          avail_mp3 > queued_flac, f"{avail_mp3:.0f} vs {queued_flac:.0f}")
    check("artist found in the folder, not the filename",
          app.match_ratio(r"Radiohead\Pablo Honey\02 Creep.flac",
                          "Radiohead", "Creep") == 1.0)
    check("a matching artist cannot rescue the wrong song",
          app.match_ratio(r"Radiohead\Pablo Honey\04 Karma Police.flac",
                          "Radiohead", "Creep") == 0.0)
    check("reordered filename still matches",
          app.match_ratio(r"x\Creep - Radiohead.mp3", "Radiohead", "Creep") == 1.0)

    print("\ncandidate ladder")
    fake = FakeSlskd(per_query={
        "Billie Jean": [peer("alice", r"m\Michael Jackson\Billie Jean.flac"),
                        peer("bob", r"n\Michael Jackson\Billie Jean.mp3"),
                        peer("carol", r"o\Michael Jackson\Billie Jean.flac")],
    })
    app.slskd_client = fake
    make_job("j1")
    asyncio.run(app.pick_and_enqueue("j1", "Michael Jackson", "Billie Jean"))
    r = job_row("j1")
    check("blocklisted artist recovered by the title rung",
          "Billie Jean" in fake.searched, fake.searched)
    check("job is downloading", r["state"] == "downloading", r["detail"])
    check("runners-up were kept", len(json.loads(r["candidates"] or "[]")) == 2,
          r["candidates"])
    check("consecutive candidates are different peers",
          len({c["username"] for c in json.loads(r["candidates"])}) == 2)

    print("\nfallback on failure")
    fake.enqueue_status = {}
    asyncio.run(app.abandon(job_row("j1"), "Transfer rejected: File not shared."))
    r = job_row("j1")
    check("a rejection advances instead of failing",
          r["state"] == "downloading", f"{r['state']} / {r['detail']}")
    check("the failed peer is recorded",
          "File not shared" in (r["tried"] or ""), r["tried"])
    check("attempt counter advanced", r["attempts"] == 2, r["attempts"])

    asyncio.run(app.abandon(job_row("j1"), "timed out"))
    asyncio.run(app.abandon(job_row("j1"), "timed out"))
    r = job_row("j1")
    check("only fails once the ladder is exhausted", r["state"] == "failed", r["detail"])
    check("failure names every attempt, not just the last",
          "tried 3 peers" in (r["detail"] or ""), r["detail"])

    print("\nattempt cap")
    fake2 = FakeSlskd(per_query={"Creep": [
        peer(f"u{i}", rf"x\Radiohead\Creep.flac") for i in range(8)]})
    # every peer refuses to accept the transfer
    fake2.enqueue_status = {f"u{i}": 500 for i in range(8)}
    app.slskd_client = fake2
    make_job("j2", "Radiohead", "Creep")
    asyncio.run(app.pick_and_enqueue("j2", "Radiohead", "Creep"))
    r = job_row("j2")
    check("gives up at MAX_ATTEMPTS, not at the end of the list",
          r["state"] == "failed" and r["attempts"] == app.MAX_ATTEMPTS,
          f"state={r['state']} attempts={r['attempts']}")
    check("unreachable peers are reported as such",
          "could not reach peer" in (r["tried"] or ""), r["detail"])

    print("\nempty results")
    fake3 = FakeSlskd(per_query={})
    app.slskd_client = fake3
    make_job("j3", "Nobody", "Nothing")
    asyncio.run(app.pick_and_enqueue("j3", "Nobody", "Nothing"))
    r = job_row("j3")
    check("says the search was empty, not that matching failed",
          "Soulseek returned nothing" in (r["detail"] or ""), r["detail"])
    check("names the queries it tried", "'Nothing'" in (r["detail"] or ""), r["detail"])

    print("\ntransient slskd")
    class Flaky(FakeSlskd):
        def __init__(self):
            super().__init__(per_query={"Creep": [peer("z", r"x\Creep.flac")]})
            self.n = 0

        async def post(self, url, json=None):
            if url == "/api/v0/searches":
                self.n += 1
                if self.n < 3:
                    raise RuntimeError("")   # httpx timeouts stringify to nothing
            return await FakeSlskd.post(self, url, json=json)

    flaky = Flaky()
    app.slskd_client = flaky
    make_job("j4", "Radiohead", "Creep")
    asyncio.run(app.pick_and_enqueue("j4", "Radiohead", "Creep"))
    check("a blip is retried rather than failing the job",
          job_row("j4")["state"] == "downloading", job_row("j4")["detail"])

    print("\nsearch hygiene")
    check("finished searches are cleaned up", len(flaky.deleted) >= 1)

    print("\nstall clock")
    from datetime import datetime, timedelta, timezone
    with app.db() as c:
        c.execute("UPDATE jobs SET started_at=? WHERE id=?",
                  ((datetime.now(timezone.utc) - timedelta(minutes=90)).isoformat(),
                   "j4"))
    check("a 90-minute-old transfer reads as stalled",
          app.stalled_for(job_row("j4")) > app.STALL_MINUTES,
          app.stalled_for(job_row("j4")))
    with app.db() as c:
        c.execute("UPDATE jobs SET started_at=NULL WHERE id=?", ("j4",))
    check("a job with no clock is not spuriously stalled",
          app.stalled_for(job_row("j4")) == 0.0)

    print(f"\n{ok} passed, {fail} failed")
    return 1 if fail else 0


if __name__ == "__main__":
    sys.exit(main())