146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
"""持久任务队列单元测试（内存 SQLite，不依赖网络）。"""
|
|
from datetime import datetime, timedelta
|
|
|
|
from sqlmodel import Session, SQLModel, create_engine, select
|
|
|
|
from app.models import TaskQueueItem
|
|
from app import task_queue as tq
|
|
|
|
|
|
def _session():
|
|
engine = create_engine("sqlite://", connect_args={"check_same_thread": False})
|
|
SQLModel.metadata.create_all(engine)
|
|
return Session(engine)
|
|
|
|
|
|
def _past(seconds=10):
|
|
return (datetime.now() - timedelta(seconds=seconds)).isoformat()
|
|
|
|
|
|
def _future(seconds=60):
|
|
return (datetime.now() + timedelta(seconds=seconds)).isoformat()
|
|
|
|
|
|
def test_enqueue_idempotent():
|
|
with _session() as s:
|
|
id1 = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={"a": 1})
|
|
id2 = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={"a": 1})
|
|
assert id1 == id2
|
|
id3 = tq.enqueue(s, TaskQueueItem, task_key="run:2", payload={"a": 2})
|
|
assert id3 != id1
|
|
|
|
|
|
def test_lease_next_respects_next_run_at():
|
|
with _session() as s:
|
|
tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
row = s.exec(select(TaskQueueItem)).first()
|
|
row.next_run_at = _future(60)
|
|
s.add(row); s.commit()
|
|
assert tq.lease_next(s, TaskQueueItem, "w1") is None
|
|
row.next_run_at = _past(1)
|
|
s.add(row); s.commit()
|
|
leased = tq.lease_next(s, TaskQueueItem, "w1")
|
|
assert leased is not None and leased.status == "running"
|
|
assert leased.attempt == 1 and leased.worker_id == "w1"
|
|
|
|
|
|
def test_lease_steals_expired_lease():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
leased = tq.lease_next(s, TaskQueueItem, "w1")
|
|
assert leased.id == task_id
|
|
row = s.get(TaskQueueItem, task_id)
|
|
row.lease_until = _past(10)
|
|
s.add(row); s.commit()
|
|
stolen = tq.lease_next(s, TaskQueueItem, "w2")
|
|
assert stolen is not None and stolen.worker_id == "w2"
|
|
assert stolen.attempt == 2
|
|
|
|
|
|
def test_heartbeat_only_by_owner():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
tq.lease_next(s, TaskQueueItem, "w1")
|
|
assert tq.heartbeat(s, TaskQueueItem, task_id, "w1") is True
|
|
assert tq.heartbeat(s, TaskQueueItem, task_id, "w2") is False
|
|
|
|
|
|
def test_fail_requeues_with_backoff_then_dead():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
tq.lease_next(s, TaskQueueItem, "w1")
|
|
status = tq.fail(s, TaskQueueItem, task_id, "w1", "第一次失败")
|
|
assert status == "pending"
|
|
row = s.get(TaskQueueItem, task_id)
|
|
assert row.attempt == 1 and row.error == "第一次失败"
|
|
assert _parse(row.next_run_at) > datetime.now()
|
|
# 退避时间清零后连续失败到上限 → 死信
|
|
for _ in range(2):
|
|
row = s.get(TaskQueueItem, task_id)
|
|
row.next_run_at = _past(1)
|
|
s.add(row); s.commit()
|
|
row2 = tq.lease_next(s, TaskQueueItem, "w1")
|
|
tq.fail(s, TaskQueueItem, row2.id, "w1", "再失败")
|
|
row = s.get(TaskQueueItem, task_id)
|
|
assert row.status == "dead"
|
|
|
|
|
|
def test_requeue_dead():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
row = s.get(TaskQueueItem, task_id)
|
|
row.status = "dead"
|
|
s.add(row); s.commit()
|
|
assert tq.requeue_dead(s, TaskQueueItem, task_id) is True
|
|
row = s.get(TaskQueueItem, task_id)
|
|
assert row.status == "pending" and row.attempt == 0
|
|
assert tq.requeue_dead(s, TaskQueueItem, task_id) is False
|
|
|
|
|
|
def test_complete_and_stats():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="run:1", payload={})
|
|
tq.lease_next(s, TaskQueueItem, "w1")
|
|
assert tq.complete(s, TaskQueueItem, task_id, "w1") is True
|
|
stats = tq.queue_stats(s, TaskQueueItem)
|
|
assert stats["total"] == 1
|
|
assert stats["by_status"]["done"] == 1
|
|
assert stats["done"] == 1
|
|
|
|
|
|
def _parse(value):
|
|
return datetime.fromisoformat(value)
|
|
|
|
|
|
def test_lease_next_filters_task_key_prefix():
|
|
with _session() as s:
|
|
tq.enqueue(s, TaskQueueItem, task_key="training-map:1", payload={})
|
|
tq.enqueue(s, TaskQueueItem, task_key="campaign-run:1", payload={})
|
|
assert tq.lease_next(s, TaskQueueItem, "dsh", task_key_prefix="campaign-%").task_key == "campaign-run:1"
|
|
|
|
|
|
def test_enqueue_updates_pending_payload_only():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="campaign-parse:1", payload={"advice": "old"})
|
|
tq.enqueue(s, TaskQueueItem, task_key="campaign-parse:1",
|
|
payload={"advice": "new"}, update_existing=True)
|
|
row = s.get(TaskQueueItem, task_id)
|
|
assert row.payload == {"advice": "new"} and row.status == "pending"
|
|
|
|
row.status = "running"
|
|
row.worker_id = "dsh1"
|
|
s.add(row); s.commit()
|
|
tq.enqueue(s, TaskQueueItem, task_key="campaign-parse:1",
|
|
payload={"advice": "ignored"}, update_existing=True)
|
|
s.refresh(row)
|
|
assert row.payload == {"advice": "new"} and row.worker_id == "dsh1"
|
|
|
|
|
|
def test_claim_by_key_pending_and_owner_only():
|
|
with _session() as s:
|
|
task_id = tq.enqueue(s, TaskQueueItem, task_key="campaign-run:9", payload={})
|
|
claimed = tq.claim_by_key(s, TaskQueueItem, "campaign-run:9", "inline-a")
|
|
assert claimed.id == task_id and claimed.worker_id == "inline-a"
|
|
assert tq.claim_by_key(s, TaskQueueItem, "campaign-run:9", "inline-b") is None
|
|
assert tq.claim_by_key(s, TaskQueueItem, "campaign-run:9", "inline-a") is not None
|