siftlode/backend/tests/test_download_limits.py

"""DB-backed tests for the two download limits:

- **#2 claim fairness** (`worker._claim_job`): users already AT their per-user concurrency limit
  are excluded IN THE SQL, so one user flooding the queue past the fetch window can't starve
  everyone behind them. The old fixed `LIMIT 50` window did exactly that.
- **#3 list cap** (`routes.downloads.list_downloads` / `shared_with_me`): the per-user listings
  are bounded so a heavy library can't return an unbounded result set.

Uses the `db` fixture (real Postgres) — skipped when no database is reachable. `_claim_job` runs in
its OWN session, so every test COMMITS its setup before calling it.
"""

from app.config import settings
from app.downloads import errors as dlerrors
from app.downloads import formats, pipeline, quota
from app.models import User
from app.downloads.models import DownloadJob, DownloadQuota, MediaAsset
from app.worker import _claim_job


def _job(db, user_id, queue_pos, status="queued", asset_id=None):
    j = DownloadJob(
        user_id=user_id,
        asset_id=asset_id,
        source_kind="youtube",
        source_ref="x" * 11,
        status=status,
        queue_pos=queue_pos,
        profile_snapshot={},
    )
    db.add(j)
    return j


def _asset(db, status, ref):
    a = MediaAsset(source_kind="url", source_ref=ref, format_sig="s", status=status, ref_count=1)
    db.add(a)
    db.flush()
    return a


def _quota(db, user_id, max_concurrent, unlimited=False):
    db.add(
        DownloadQuota(
            user_id=user_id,
            max_bytes=10**12,
            max_concurrent=max_concurrent,
            max_jobs=10**6,
            unlimited=unlimited,
        )
    )


# --- #2 · claim fairness --------------------------------------------------------------------


def test_capped_user_flood_does_not_starve_users_behind_the_window(db):
    a = User(email="flooder@example.com")
    b = User(email="waiter@example.com")
    db.add_all([a, b])
    db.flush()
    _quota(
        db, a.id, max_concurrent=1
    )  # A: 1 concurrent, already running one -> all A's queued are ineligible
    _quota(db, b.id, max_concurrent=3)
    _job(db, a.id, queue_pos=0, status="running")
    for i in range(1, 61):  # A floods the FRONT with far more than the 50-row fetch window
        _job(db, a.id, queue_pos=i)
    bjob = _job(db, b.id, queue_pos=100)  # B sits behind the flood
    db.commit()

    assert _claim_job() == bjob.id  # B is served despite A's 60-job flood filling the old window
    db.expire_all()
    assert db.get(DownloadJob, bjob.id).status == "running"


def test_eligible_user_claims_oldest_by_queue_pos(db):
    u = User(email="solo@example.com")
    db.add(u)
    db.flush()
    _quota(db, u.id, max_concurrent=2)
    first = _job(db, u.id, queue_pos=1)
    _job(db, u.id, queue_pos=2)
    db.commit()
    assert _claim_job() == first.id


def test_all_candidates_capped_returns_none(db):
    u = User(email="capped@example.com")
    db.add(u)
    db.flush()
    _quota(db, u.id, max_concurrent=1)
    _job(db, u.id, queue_pos=0, status="running")  # at the cap
    _job(db, u.id, queue_pos=1)  # queued but ineligible
    db.commit()
    assert _claim_job() is None


def test_stored_zero_concurrency_is_floored_not_starving(db):
    # A stray max_concurrent=0 is only reachable by a direct DB edit (admin_set_quota floors at 1),
    # but it must not wedge the claim loop. The effective limit floors to 1 on BOTH the read path
    # (resolve -> at_concurrency_limit) and in the claim SQL (GREATEST(..., 1)), so the two "at
    # limit" checks agree: the user's own first job is claimable and its flood never monopolizes
    # the window with permanently-ineligible jobs. (Before the floor, _claim_job returned None here
    # — the guard rejected all 60 of A's jobs — and any user behind them starved.)
    a = User(email="zero@example.com")
    b = User(email="behind@example.com")
    db.add_all([a, b])
    db.flush()
    _quota(db, a.id, max_concurrent=0)  # nonsensical stored value
    _quota(db, b.id, max_concurrent=2)
    for i in range(60):
        _job(db, a.id, queue_pos=i)
    _job(db, b.id, queue_pos=100)
    db.commit()

    claimed = _claim_job()
    assert claimed is not None  # the loop is not wedged
    db.expire_all()
    assert db.get(DownloadJob, claimed).user_id == a.id  # A's own first job runs (floored to 1)


def test_job_whose_asset_is_already_downloading_is_skipped(db):
    """B-sub-3: a queued job whose shared asset is ALREADY `downloading` (another worker owns it) must
    not be claimed — otherwise a second worker claims it, finds the asset busy, requeues, and ping-
    pongs on it, starving everyone behind. The claim SQL excludes it, so the job behind it (on a
    different, free asset) is served instead."""
    u = User(email="asset-excl@example.com")
    db.add(u)
    db.flush()
    _quota(
        db, u.id, max_concurrent=5
    )  # high → the ASSET rule, not the per-user rule, does the excluding
    busy = _asset(db, status="downloading", ref="busy")
    free = _asset(db, status="pending", ref="free")
    _job(
        db, u.id, queue_pos=1, status="running", asset_id=busy.id
    )  # the holder actually downloading it
    competitor = _job(db, u.id, queue_pos=2, asset_id=busy.id)  # same asset — must be skipped
    behind = _job(db, u.id, queue_pos=3, asset_id=free.id)  # different asset — should be served
    db.commit()

    assert _claim_job() == behind.id  # the competitor on the busy asset was skipped
    db.expire_all()
    assert db.get(DownloadJob, competitor.id).status == "queued"  # still queued, never claimed


def test_unlimited_user_is_never_excluded(db):
    u = User(email="unlimited@example.com")
    db.add(u)
    db.flush()
    _quota(db, u.id, max_concurrent=1, unlimited=True)
    for i in range(3):  # 3 running, well past max_concurrent — but unlimited bypasses the cap
        _job(db, u.id, queue_pos=i, status="running")
    q = _job(db, u.id, queue_pos=10)
    db.commit()
    assert _claim_job() == q.id


def test_a_suspended_users_queued_jobs_wait(db):
    # Suspending a user blocks their downloads (quota.resolve's advice): a job they queued before
    # stays queued, unclaimed, and runs once the suspension is lifted.
    s = User(email="suspended@example.com", is_suspended=True)
    o = User(email="other@example.com")
    db.add_all([s, o])
    db.flush()
    waits = _job(db, s.id, queue_pos=1)
    served = _job(db, o.id, queue_pos=2)
    db.commit()
    assert _claim_job() == served.id
    assert _claim_job() is None
    db.expire_all()
    assert db.get(DownloadJob, waits.id).status == "queued"
    db.get(User, s.id).is_suspended = False
    db.commit()
    assert _claim_job() == waits.id


# --- #3 · list cap --------------------------------------------------------------------------


def test_list_downloads_is_capped(db, monkeypatch):
    from app.downloads import queries
    from app.routes import downloads as dl

    u = User(email="heavy@example.com")
    db.add(u)
    db.flush()
    for i in range(3):
        _job(db, u.id, queue_pos=i, status="done")
    db.commit()

    monkeypatch.setattr(queries, "MAX_LIST", 2)
    rows = dl.list_downloads(user=u, db=db)
    assert len(rows) == 2  # bounded by the cap, not the 3 rows on disk


# --- the worker-side max_bytes remainder ----------------------------------------------------------


def _sized_asset(db, size_bytes, ref):
    a = MediaAsset(
        source_kind="url",
        source_ref=ref,
        format_sig="s",
        status="ready",
        ref_count=1,
        size_bytes=size_bytes,
    )
    db.add(a)
    db.flush()
    return a


def test_remaining_bytes_is_none_for_an_unlimited_user(db):
    u = User(email="unlimited@example.com")
    db.add(u)
    db.flush()
    _quota(db, u.id, max_concurrent=2, unlimited=True)
    db.flush()
    # None, not a number: "no cap" and "no budget left" must never be the same value downstream —
    # build_ydl_opts leaves max_filesize unset for both, so only the caller can tell them apart.
    assert quota.remaining_bytes(db, u.id) is None


def test_remaining_bytes_is_the_unspent_part_of_the_cap(db):
    u = User(email="partial@example.com")
    db.add(u)
    db.flush()
    db.add(
        DownloadQuota(user_id=u.id, max_bytes=1000, max_concurrent=2, max_jobs=100, unlimited=False)
    )
    a = _sized_asset(db, 400, "partial-1")
    _job(db, u.id, queue_pos=0, status="done", asset_id=a.id)
    db.flush()
    assert quota.remaining_bytes(db, u.id) == 600


def test_remaining_bytes_is_zero_once_the_cap_is_met(db):
    # THE case this exists for: `check_enqueue` let these jobs in while the user was under the cap,
    # and an earlier one has filled it since. Zero is what makes the worker refuse; if this ever
    # returned the raw negative or None instead, the user already over quota would be the only one
    # downloading with no size ceiling at all.
    u = User(email="atcap@example.com")
    db.add(u)
    db.flush()
    db.add(
        DownloadQuota(user_id=u.id, max_bytes=1000, max_concurrent=2, max_jobs=100, unlimited=False)
    )
    a = _sized_asset(db, 1500, "atcap-1")  # overshot: one job finished larger than the remainder
    _job(db, u.id, queue_pos=0, status="done", asset_id=a.id)
    db.flush()
    assert quota.remaining_bytes(db, u.id) == 0


def test_a_refusal_for_one_user_does_not_poison_the_shared_asset(db, monkeypatch, tmp_path):
    """A user at their cap fails their own job, and another user's queued job for the same video
    still downloads.

    MediaAssets are deduped across users, and the generic failure handler writes the message onto
    that shared row — correct for an upstream failure, wrong for "YOUR storage is full", which
    `worker._process_job` would then hand verbatim to someone with an empty library.
    """
    monkeypatch.setattr(settings, "files_root", str(tmp_path))
    over, other = User(email="over@example.com"), User(email="other@example.com")
    db.add_all([over, other])
    db.flush()
    db.add(
        DownloadQuota(
            user_id=over.id, max_bytes=1000, max_concurrent=2, max_jobs=100, unlimited=False
        )
    )
    full = _sized_asset(db, 1500, "already-held")
    _job(db, over.id, queue_pos=0, status="done", asset_id=full.id)

    # `downloading` is the state the worker leaves it in — `_claim_asset` compare-and-sets it before
    # the child runs — and `_reset_asset_pending` only acts on that status. Seeded as `pending` this
    # assertion would pass with the release deleted.
    shared = MediaAsset(
        source_kind="youtube",
        source_ref="x" * 11,
        format_sig="s",
        status="downloading",
        ref_count=2,
    )
    db.add(shared)
    db.flush()
    mine = _job(db, over.id, queue_pos=1, status="running", asset_id=shared.id)
    theirs = _job(db, other.id, queue_pos=2, asset_id=shared.id)
    db.commit()

    # The refusal happens before extraction, so nothing needs stubbing — if it ever stops happening
    # first, this test starts reaching the network, which is its own signal.
    pipeline._download(mine.id, shared.id, "youtube", "x" * 11, {}, over.id)

    db.expire_all()
    assert db.get(DownloadJob, mine.id).status == "error"
    assert db.get(DownloadJob, mine.id).error_code == "quota_exceeded"
    # The asset is released, NOT marked error — otherwise the next line's job inherits the refusal.
    assert db.get(MediaAsset, shared.id).status == "pending"
    assert db.get(MediaAsset, shared.id).error is None
    assert db.get(DownloadJob, theirs.id).status == "queued"


def test_a_download_over_the_space_left_says_so_and_releases_the_shared_asset(
    db, monkeypatch, tmp_path
):
    # yt-dlp refuses a file over `max_filesize` by saying so and returning with no file (pinned in
    # test_formats). That is the owner's storage, not a processing fault: their refusal, in words
    # true of a cap with space still left, and the shared asset goes back for the next job.
    monkeypatch.setattr(settings, "files_root", str(tmp_path))
    u = User(email="nearcap@example.com")
    db.add(u)
    db.flush()
    db.add(
        DownloadQuota(user_id=u.id, max_bytes=1000, max_concurrent=2, max_jobs=100, unlimited=False)
    )
    held = _sized_asset(db, 400, "near-held")
    _job(db, u.id, queue_pos=0, status="done", asset_id=held.id)
    shared = MediaAsset(
        source_kind="youtube",
        source_ref="y" * 11,
        format_sig="s",
        status="downloading",
        ref_count=1,
    )
    db.add(shared)
    db.flush()
    mine = _job(db, u.id, queue_pos=1, status="running", asset_id=shared.id)
    db.commit()
    caps = []

    def refuses_over_the_cap(*args, max_bytes=None, **kwargs):
        caps.append(max_bytes)
        formats._YDL_LOG.debug(
            f"[download] File is larger than max-filesize (2000 bytes > {max_bytes} bytes). Aborting."
        )
        return {"id": "y" * 11, "requested_downloads": [{}]}

    monkeypatch.setattr(pipeline, "_extract_with_fallback", refuses_over_the_cap)
    pipeline._download(mine.id, shared.id, "youtube", "y" * 11, {}, u.id)

    db.expire_all()
    job = db.get(DownloadJob, mine.id)
    assert caps == [600]
    assert (job.status, job.error_code) == ("error", "quota_exceeded")
    assert job.error == dlerrors.DOWNLOAD_TOO_BIG_MESSAGE
    assert db.get(MediaAsset, shared.id).status == "pending"


def test_a_download_that_outgrows_the_space_left_unannounced_is_refused_and_not_placed(
    db, monkeypatch, tmp_path
):
    # yt-dlp's `max_filesize` only stops a plain HTTP file that announces its size: YouTube's
    # fragmented formats never meet it (measured on localdev: 718 MB fetched with 1 MB left). The
    # file it produced is held to what is left by its real size, as a clip is.
    monkeypatch.setattr(settings, "files_root", str(tmp_path))
    u = User(email="unannounced@example.com")
    db.add(u)
    db.flush()
    db.add(
        DownloadQuota(user_id=u.id, max_bytes=1000, max_concurrent=2, max_jobs=100, unlimited=False)
    )
    held = _sized_asset(db, 400, "unannounced-held")
    _job(db, u.id, queue_pos=0, status="done", asset_id=held.id)
    shared = MediaAsset(
        source_kind="youtube",
        source_ref="z" * 11,
        format_sig="s",
        status="downloading",
        ref_count=1,
    )
    db.add(shared)
    db.flush()
    mine = _job(db, u.id, queue_pos=1, status="running", asset_id=shared.id)
    db.commit()

    def fetched_past_the_cap(job_id, asset_id, url, spec, staging, adapter, **kw):
        f = staging / "zzzzzzzzzzz.mp4"
        f.write_bytes(b"v" * 2000)  # 600 left
        return {"id": "z" * 11, "requested_downloads": [{"filepath": str(f)}]}

    monkeypatch.setattr(pipeline, "_extract_with_fallback", fetched_past_the_cap)
    monkeypatch.setattr(pipeline, "_ensure_phone_safe", lambda _j, produced, *_a: produced)
    pipeline._download(mine.id, shared.id, "youtube", "z" * 11, {}, u.id)

    db.expire_all()
    job = db.get(DownloadJob, mine.id)
    assert (job.status, job.error_code, job.error) == (
        "error",
        "quota_exceeded",
        dlerrors.DOWNLOAD_TOO_BIG_MESSAGE,
    )
    asset = db.get(MediaAsset, shared.id)
    assert (asset.status, asset.rel_path) == ("pending", None)
    assert not any(p.suffix == ".mp4" for p in tmp_path.rglob("*") if "staging" not in p.parts)