67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
"""调度服务单元测试。"""
|
||
import pytest
|
||
|
||
from app.constants import 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
|
||
|
||
|
||
@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 = Scheduler(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_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 = Scheduler(redis, task_repo, agent_repo)
|
||
n = await sched.dispatch_pending(10)
|
||
assert n == 2 |