174 lines
7.5 KiB
Python
174 lines
7.5 KiB
Python
"""调度服务单元测试。"""
|
||
import pytest
|
||
|
||
from app.constants import AgentStatus, TaskStatus
|
||
from app.models.schemas import AgentInfo, TaskInfo
|
||
from app.repository.agent_repo import AgentRepo
|
||
from app.repository.task_repo import TaskRepo
|
||
from app.services.scheduler import Scheduler
|
||
|
||
|
||
class _StubRelay:
|
||
"""模拟 relay：不发起真实网络请求，直接返回推送成功。"""
|
||
|
||
def __init__(self, ok: bool = True):
|
||
self.ok = ok
|
||
|
||
async def dispatch_command(self, agent_id: str, request_id: str, payload: dict) -> bool:
|
||
return self.ok
|
||
|
||
|
||
def _sched(redis, task_repo, agent_repo):
|
||
return Scheduler(redis, task_repo, agent_repo, relay=_StubRelay())
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_binds_min_load_agent(redis):
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
|
||
# 两个 compile Agent，b2 负载更低
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=1))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=0))
|
||
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
|
||
sched = _sched(redis, task_repo, agent_repo)
|
||
ok = await sched.dispatch("t1")
|
||
assert ok is True
|
||
|
||
task = await task_repo.get("t1")
|
||
assert task.agent_id == "b2"
|
||
assert task.status == TaskStatus.RUNNING
|
||
assert "t1" in await task_repo.running_tasks()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_pick_candidate_lowest_priority(redis):
|
||
agent_repo = AgentRepo(redis)
|
||
|
||
# 三个 compile Agent：b1 负载最低但优先级高(最不被优先)，b2 负载高但优先级低(最优先)
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=5))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=1, priority=1))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=3))
|
||
|
||
sched = Scheduler(redis)
|
||
picked = await sched.pick_candidate(["compile"])
|
||
assert picked is not None
|
||
assert picked.agent_id == "b2" # 优先级最小优先分配
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_pick_candidate_same_priority_min_load(redis):
|
||
agent_repo = AgentRepo(redis)
|
||
|
||
# 同优先级(默认3)时，负载更低者优先
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=1))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=0))
|
||
|
||
sched = Scheduler(redis)
|
||
picked = await sched.pick_candidate(["compile"])
|
||
assert picked is not None
|
||
assert picked.agent_id == "b2"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_no_agent_stays_pending(redis):
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
|
||
sched = Scheduler(redis, task_repo, agent_repo)
|
||
ok = await sched.dispatch("t1")
|
||
assert ok is False
|
||
task = await task_repo.get("t1")
|
||
assert task.status == TaskStatus.PENDING
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_skips_full_agent(redis):
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
# 唯一 agent 已满载
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=1, current_load=1))
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
|
||
sched = Scheduler(redis, task_repo, agent_repo)
|
||
assert await sched.dispatch("t1") is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_pending(redis):
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
await task_repo.create(TaskInfo(request_id="t2", task_type="compile", task_tags=["compile"]))
|
||
|
||
sched = _sched(redis, task_repo, agent_repo)
|
||
n = await sched.dispatch_pending(10)
|
||
assert n == 2
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_excludes_offline_unavailable_stopping(redis):
|
||
"""调度排除 offline（心跳失联）、unavailable（手动不可用）与 stopping 的 agent。"""
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
# b1 手动不可用，b2 正在停止，b3 离线，b4 就绪
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2))
|
||
await agent_repo.upsert(AgentInfo(agent_id="b4", endpoint="e4", agent_tags=["compile"], max_concurrent=2))
|
||
await agent_repo.set_unavailable("b1")
|
||
await agent_repo.update_status("b2", AgentStatus.STOPPING)
|
||
await agent_repo.update_status("b3", AgentStatus.OFFLINE)
|
||
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
sched = _sched(redis, task_repo, agent_repo)
|
||
ok = await sched.dispatch("t1")
|
||
assert ok is True
|
||
task = await task_repo.get("t1")
|
||
assert task.agent_id == "b4"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_processing_agent_with_capacity(redis):
|
||
"""processing 的 agent 若还有并发容量，仍可接收新任务。"""
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
|
||
await agent_repo.update_status("b1", AgentStatus.PROCESSING)
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
sched = _sched(redis, task_repo, agent_repo)
|
||
assert await sched.dispatch("t1") is True
|
||
assert (await task_repo.get("t1")).agent_id == "b1"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_marks_agent_processing(redis):
|
||
"""调度下发成功后，agent 状态流转为 processing。"""
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
sched = _sched(redis, task_repo, agent_repo)
|
||
assert await sched.dispatch("t1") is True
|
||
assert (await agent_repo.get("b1")).status == AgentStatus.PROCESSING
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dispatch_push_fail_reverts_agent_to_ready(redis):
|
||
"""推送失败时任务回退 pending，agent 释放负载并回到 ready。"""
|
||
task_repo = TaskRepo(redis)
|
||
agent_repo = AgentRepo(redis)
|
||
await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
|
||
await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
|
||
sched = Scheduler(redis, task_repo, agent_repo, relay=_StubRelay(ok=False))
|
||
assert await sched.dispatch("t1") is False
|
||
task = await task_repo.get("t1")
|
||
assert task.status == TaskStatus.PENDING
|
||
assert not task.agent_id
|
||
a = await agent_repo.get("b1")
|
||
assert a.status == AgentStatus.READY
|
||
assert a.current_load == 0 |