Lyra/worker/tests/test_slskd_cancel.py

"""slskd download behaviour: abandoned transfers are cancelled (no pile-up), a stalled peer is
dropped fast, a slow-but-progressing peer is kept, and search candidates are ranked fast/free
peers first."""
import json

import pytest

import lyra_worker.adapters._slskd as slskd_mod
from lyra_worker.adapters._slskd import (
    _ETA_CAP_SECONDS,
    _MAX_XFER_POLLS,
    _MIN_XFER_POLLS,
    _POLL_SECONDS,
    SlskdClient,
    _eta_seconds,
    _max_polls,
    _parse_search_responses,
)


class _Resp:
    def __init__(self, payload=None):
        self._p = payload if payload is not None else {}

    def raise_for_status(self):
        pass

    def json(self):
        return self._p


class _SeqRequests:
    """post OK; get() returns successive poll payloads (last one repeats); records DELETEs."""
    def __init__(self, polls):
        self._polls = list(polls)
        self._i = 0
        self.deleted = []

    def post(self, url, **kw):
        return _Resp({})

    def get(self, url, **kw):
        p = self._polls[min(self._i, len(self._polls) - 1)]
        self._i += 1
        return _Resp(p)

    def delete(self, url, **kw):
        self.deleted.append(url)
        return _Resp({})


def _poll(state, bytes_transferred, fid="t1"):
    return {"directories": [{"files": [
        {"filename": "f1.flac", "size": 100, "state": state,
         "bytesTransferred": bytes_transferred, "id": fid},
    ]}]}


def _client(tmp_path):
    return SlskdClient({"slskd.url": "http://x", "slskd.api_key": "k",
                        "slskd.downloadsDir": str(tmp_path)})


_REF = json.dumps({"username": "peerA", "files": [{"filename": "f1.flac", "size": 100}]})


def test_stalled_peer_is_abandoned_and_cancelled(tmp_path, monkeypatch):
    # No byte progress across polls → stall → TimeoutError → the transfer is cancelled.
    monkeypatch.setattr(slskd_mod, "requests", _SeqRequests([_poll("InProgress", 0)]))
    monkeypatch.setattr(slskd_mod.time, "sleep", lambda _s: None)
    monkeypatch.setattr(slskd_mod, "_STALL_POLLS", 3)
    fake = slskd_mod.requests

    with pytest.raises(TimeoutError):
        _client(tmp_path).download(_REF, str(tmp_path / "dest"), lambda *_a: None)
    assert any(u.endswith("/api/v0/transfers/downloads/peerA/t1") for u in fake.deleted)


def test_error_state_is_cancelled(tmp_path, monkeypatch):
    monkeypatch.setattr(slskd_mod, "requests", _SeqRequests([_poll("Completed, Errored", 0)]))
    monkeypatch.setattr(slskd_mod.time, "sleep", lambda _s: None)
    fake = slskd_mod.requests
    with pytest.raises(RuntimeError):
        _client(tmp_path).download(_REF, str(tmp_path / "dest"), lambda *_a: None)
    assert any(u.endswith("/peerA/t1") for u in fake.deleted)


def test_slow_but_progressing_peer_is_not_abandoned(tmp_path, monkeypatch):
    # Bytes climb for more polls than _STALL_POLLS, then complete → success, not a premature stall.
    (tmp_path / "01.flac").write_bytes(b"audio")  # so _retrieve finds the finished file
    ref = json.dumps({"username": "peerA", "files": [{"filename": r"d\Album\01.flac", "size": 5}]})
    polls = [
        {"directories": [{"files": [{"filename": r"d\Album\01.flac", "size": 5,
                                     "state": "InProgress", "bytesTransferred": b, "id": "t1"}]}]}
        for b in (10, 20, 30, 40, 50)
    ] + [
        {"directories": [{"files": [{"filename": r"d\Album\01.flac", "size": 5,
                                     "state": "Completed, Succeeded", "bytesTransferred": 5, "id": "t1"}]}]}
    ]
    monkeypatch.setattr(slskd_mod, "requests", _SeqRequests(polls))
    monkeypatch.setattr(slskd_mod.time, "sleep", lambda _s: None)
    monkeypatch.setattr(slskd_mod, "_STALL_POLLS", 2)  # would fire early if progress were ignored

    out = _client(tmp_path).download(ref, str(tmp_path / "dest"), lambda *_a: None)
    assert out["track_count"] == 1
    assert slskd_mod.requests.deleted == []  # completed cleanly → nothing cancelled


def test_parse_ranks_free_and_fast_peers_first():
    responses = [
        {"username": "slow", "hasFreeUploadSlot": True, "uploadSpeed": 100, "queueLength": 0,
         "files": [{"filename": r"a\Artist - Album\01.flac", "size": 5}]},
        {"username": "fast", "hasFreeUploadSlot": True, "uploadSpeed": 9000, "queueLength": 0,
         "files": [{"filename": r"b\Artist - Album\01.flac", "size": 5}]},
        {"username": "queued", "hasFreeUploadSlot": False, "uploadSpeed": 99999, "queueLength": 40,
         "files": [{"filename": r"c\Artist - Album\01.flac", "size": 5}]},
    ]
    peers = [json.loads(c["source_ref"])["username"] for c in _parse_search_responses(responses)]
    assert peers == ["fast", "slow", "queued"]  # free+fast first; no-slot peer last despite speed


def test_parse_uses_requested_artist_when_present_in_path():
    # Real "Rated R" MP3-over-FLAC bug: the artist is in a PARENT folder ("Rihanna\(2009) Rated
    # R\…"), so the "Artist - Album" folder parse yields no artist and confidence tanks. When the
    # requested artist is in the path, trust it so the candidate scores as Rihanna.
    responses = [
        {"username": "p", "hasFreeUploadSlot": True, "uploadSpeed": 1000, "queueLength": 0,
         "files": [{"filename": r"@@x\Music\Rihanna\(2009) Rated R\01. Mad House.flac", "size": 5}]},
    ]
    c = _parse_search_responses(responses, artist="Rihanna")[0]
    assert c["artist"] == "Rihanna"


def test_parse_without_artist_falls_back_to_folder_parse():
    responses = [
        {"username": "p", "hasFreeUploadSlot": True, "uploadSpeed": 1000, "queueLength": 0,
         "files": [{"filename": r"x\Some Band - Some Album\01.flac", "size": 5}]},
    ]
    c = _parse_search_responses(responses)[0]  # no artist arg → folder parse
    assert c["artist"] == "Some Band"


def test_parse_prefers_shorter_transfer_over_raw_speed():
    # The Hybrid Theory case: a peer advertising higher speed but serving a ~2x-larger rip (24-bit
    # / bloated FLAC) should lose to a peer with lower speed but a smaller standard rip that
    # finishes sooner — both are the same quality class to Lyra, so faster-to-finish wins.
    responses = [
        {"username": "bloated", "hasFreeUploadSlot": True, "uploadSpeed": 20000, "queueLength": 0,
         "files": [{"filename": rf"x\LP - Album\{i:02}.flac", "size": 40_000_000} for i in range(12)]},
        {"username": "lean", "hasFreeUploadSlot": True, "uploadSpeed": 15000, "queueLength": 0,
         "files": [{"filename": rf"y\LP - Album\{i:02}.flac", "size": 24_000_000} for i in range(12)]},
    ]
    # bloated est = 480MB/20000 = 24000s; lean est = 288MB/15000 = 19200s → lean finishes first
    peers = [json.loads(c["source_ref"])["username"] for c in _parse_search_responses(responses)]
    assert peers[0] == "lean"


def test_eta_seconds_from_measured_rate():
    assert _eta_seconds(1000, 250, 250.0) == 3  # 750 bytes left / 250 B/s
    assert _eta_seconds(1000, 0, 100.0) == 10


def test_eta_seconds_unknown_cases():
    assert _eta_seconds(1000, 250, 0.0) is None      # no speed sample yet
    assert _eta_seconds(0, 0, 100.0) is None          # unknown total size
    assert _eta_seconds(1000, 1000, 100.0) is None    # already complete


def test_max_polls_scales_backstop_to_album_size():
    # A 521MB album (Drake "Scorpion", 25 FLACs) from a slow serial peer must get far more than
    # the old flat 1200-poll/60min budget, or it gets abandoned mid-download and restarts from zero.
    scorpion_bytes = 521_000_000
    polls = _max_polls(scorpion_bytes)
    assert polls > 1200                              # more headroom than the old flat cap
    assert polls * _POLL_SECONDS >= 90 * 60          # at least ~90 min of wall-clock budget

    # A tiny single stays on the floor, not zero.
    assert _max_polls(3_000_000) == _MIN_XFER_POLLS
    assert _max_polls(0) == _MIN_XFER_POLLS

    # An enormous set is clamped to the ceiling (bounds worst-case slot hold).
    assert _max_polls(50_000_000_000) == _MAX_XFER_POLLS

    # Monotonic: bigger album never gets a smaller budget.
    assert _max_polls(200_000_000) <= _max_polls(400_000_000)


def test_eta_seconds_capped_when_peer_nearly_stalls():
    # A near-stalled peer (tiny speed) would yield an ETA far past int4, overflowing the
    # downloadEtaSeconds column write. The estimate must be clamped to a safe ceiling.
    eta = _eta_seconds(30_000_000, 0, 0.001)  # 30MB / 0.001 B/s ~= 3e10s uncapped
    assert eta == _ETA_CAP_SECONDS
    assert eta < 2_147_483_647  # fits in int4


def test_download_reports_shrinking_eta(tmp_path, monkeypatch):
    # A 900-byte file arriving 300 bytes per 3s poll → measured speed 100 B/s → ETA counts down
    # (6s, 3s) then None on completion. on_progress gets (fraction, eta_seconds).
    (tmp_path / "01.flac").write_bytes(b"x" * 900)  # so _retrieve finds the finished file
    ref = json.dumps({"username": "peerA", "files": [{"filename": r"d\Album\01.flac", "size": 900}]})
    polls = [
        {"directories": [{"files": [{"filename": r"d\Album\01.flac", "size": 900,
                                     "state": "InProgress", "bytesTransferred": b, "id": "t1"}]}]}
        for b in (300, 600)
    ] + [
        {"directories": [{"files": [{"filename": r"d\Album\01.flac", "size": 900,
                                     "state": "Completed, Succeeded", "bytesTransferred": 900, "id": "t1"}]}]}
    ]
    monkeypatch.setattr(slskd_mod, "requests", _SeqRequests(polls))
    monkeypatch.setattr(slskd_mod.time, "sleep", lambda _s: None)
    clock = iter(float(i) * 3 for i in range(100))  # 0,3,6,… deterministic monotonic time
    monkeypatch.setattr(slskd_mod.time, "monotonic", lambda: next(clock))

    calls: list = []
    _client(tmp_path).download(ref, str(tmp_path / "dest"), lambda *a: calls.append(a))

    etas = [a[1] for a in calls if len(a) > 1]
    assert etas == [6, 3, None]  # counts down at 100 B/s, then None once complete