tests/test_paid_queue.py (4929 bytes)
1 """Burst workers take paid jobs only. 2 3 Each server job goes to the paid queue or to the free queue. The web server's 4 own worker takes both, paid jobs first. A burst worker takes paid jobs only, so 5 free credits never start a machine that costs money. 6 """ 7 8 from __future__ import annotations 9 10 import time 11 from datetime import date 12 13 import fakeredis 14 import pytest 15 from rq import Queue 16 17 import worker 18 from backend import api, billing, main 19 from backend import queue as jobqueue 20 from backend.db import Job, JobStatus, SessionLocal 21 from backend.settings import settings 22 23 from .conftest import account_id, checkout_event, make_job, post_webhook 24 25 26 @pytest.fixture 27 def cloud(monkeypatch): 28 monkeypatch.setattr(settings, "cloud_enabled", True) 29 30 31 @pytest.fixture 32 def sent(monkeypatch): 33 """What the API put on the queue: (job id, paid).""" 34 calls: list[tuple[str, bool]] = [] 35 monkeypatch.setattr(api.queue, "enqueue", lambda job_id, paid=True: calls.append((job_id, paid))) 36 return calls 37 38 39 def start(client, job_id: str): 40 return client.post(f"/api/jobs/{job_id}/start", json={"language": "en"}) 41 42 43 @pytest.mark.parametrize("fields, paid", [ 44 ({"credit_spent": 1}, True), # a bought credit 45 ({"plan_credit_spent": 1}, True), # a book of a monthly plan 46 ({}, True), # an unlimited plan spends no credit 47 ({"free_credit_spent": 1}, False), 48 ({"daily_credit_on": date(2026, 9, 22)}, False), 49 ({"local": 1}, False), # a browser job never goes to a worker 50 ]) 51 def test_is_paid(fields, paid): 52 job = Job(**{"local": 0, "credit_spent": 0, "plan_credit_spent": 0, 53 "free_credit_spent": 0, "daily_credit_on": None, **fields}) 54 assert billing.is_paid(job) is paid 55 56 57 def test_a_free_credit_job_goes_to_the_free_queue(client, cloud, sent): 58 job_id = make_job(client, with_files=True) 59 assert start(client, job_id).status_code == 200 60 assert sent == [(job_id, False)] 61 62 63 def test_a_bought_credit_job_goes_to_the_paid_queue(client, cloud, sent): 64 assert start(client, make_job(client, with_files=True)).status_code == 200 # the free one 65 event = checkout_event(f"cs_{time.time_ns()}", account_id(client), "pack10", 66 email=f"paid-{time.time_ns()}@example.com") 67 assert post_webhook(client, event).status_code == 200 68 job_id = make_job(client, with_files=True) 69 assert start(client, job_id).status_code == 200 70 assert sent[-1] == (job_id, True) 71 72 73 def test_the_redis_queue_keeps_paid_and_free_jobs_apart(monkeypatch): 74 conn = fakeredis.FakeRedis() 75 monkeypatch.setattr("redis.Redis.from_url", lambda url: conn) 76 q = jobqueue.RedisQueue("redis://unused") 77 q.enqueue("job_paid", paid=True) 78 q.enqueue("job_free", paid=False) 79 assert [j.args for j in Queue(jobqueue.PAID_QUEUE, connection=conn).jobs] == [("job_paid",)] 80 assert [j.args for j in Queue(jobqueue.FREE_QUEUE, connection=conn).jobs] == [("job_free",)] 81 assert q.depth() == 2 82 83 84 def test_a_job_whose_worker_dies_gets_one_more_attempt(monkeypatch): 85 # A spot machine can vanish during a job. rq puts such a job back on its 86 # queue while the job has retries left: one retry, so two attempts at most. 87 conn = fakeredis.FakeRedis() 88 monkeypatch.setattr("redis.Redis.from_url", lambda url: conn) 89 jobqueue.RedisQueue("redis://unused").enqueue("job_1", paid=True) 90 assert Queue(jobqueue.PAID_QUEUE, connection=conn).jobs[0].retries_left == 1 91 92 93 def test_the_web_server_worker_takes_paid_jobs_first(monkeypatch): 94 monkeypatch.delenv("SUBPLZ_WEB_WORKER_QUEUES", raising=False) 95 assert worker.queue_names() == [jobqueue.PAID_QUEUE, jobqueue.FREE_QUEUE] 96 97 98 def test_a_burst_worker_takes_paid_jobs_only(monkeypatch): 99 monkeypatch.setenv("SUBPLZ_WEB_WORKER_QUEUES", "paid") 100 assert worker.queue_names() == [jobqueue.PAID_QUEUE] 101 102 103 def test_an_unknown_queue_is_refused(monkeypatch): 104 monkeypatch.setenv("SUBPLZ_WEB_WORKER_QUEUES", "paid,vip") 105 with pytest.raises(ValueError, match="vip"): 106 worker.queue_names() 107 108 109 def _running_job(client) -> str: 110 job_id = make_job(client, JobStatus.running) 111 with SessionLocal() as s: 112 s.get(Job, job_id).credit_spent = 1 113 s.commit() 114 return job_id 115 116 117 def test_a_restart_with_the_redis_queue_leaves_running_jobs_alone(client, sent, monkeypatch): 118 # With Redis, a worker on another machine still runs the job. Putting it on 119 # the queue again would convert the book twice. 120 job_id = _running_job(client) 121 monkeypatch.setattr(settings, "queue_backend", "redis") 122 main._requeue_interrupted() 123 assert sent == [] 124 with SessionLocal() as s: 125 assert s.get(Job, job_id).status == JobStatus.running 126 127 128 def test_a_restart_with_the_memory_queue_runs_interrupted_jobs_again(client, sent, monkeypatch): 129 job_id = _running_job(client) 130 monkeypatch.setattr(settings, "queue_backend", "memory") 131 main._requeue_interrupted() 132 assert (job_id, True) in sent