feat(server): add workbench API, DSH planner worker and RAG pipeline

server/tests/test_task_queue.py

@@ -0,0 +1,145 @@
+"""持久任务队列单元测试（内存 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