161 lines
6.9 KiB
Python
161 lines
6.9 KiB
Python
"""Agent 池存储层单元测试。"""
|
||
import time
|
||
|
||
import pytest
|
||
|
||
from app.constants import AgentStatus
|
||
from app.models.schemas import AgentInfo
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_upsert_and_get(agent_repo):
|
||
a = AgentInfo(agent_id="a1", endpoint="agent-1:5000", agent_tags=["compile", "test"], max_concurrent=2)
|
||
created = await agent_repo.upsert(a)
|
||
assert created is True
|
||
|
||
got = await agent_repo.get("a1")
|
||
assert got is not None
|
||
assert got.agent_tags == ["compile", "test"]
|
||
assert got.status == AgentStatus.READY
|
||
assert got.priority == 3 # 默认优先级
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_priority_roundtrip_and_update(agent_repo):
|
||
a = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], priority=1)
|
||
await agent_repo.upsert(a)
|
||
assert (await agent_repo.get("a1")).priority == 1
|
||
|
||
assert await agent_repo.update_priority("a1", 5) is True
|
||
assert (await agent_repo.get("a1")).priority == 5
|
||
|
||
assert await agent_repo.update_priority("missing", 2) is False
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_beat_and_stale(agent_repo):
|
||
a = AgentInfo(agent_id="a1", endpoint="agent-1:5000", agent_tags=["compile"])
|
||
await agent_repo.upsert(a)
|
||
assert await agent_repo.beat("a1", 1) is True
|
||
|
||
# 心跳为最新，不应 stale
|
||
stale = await agent_repo.stale_agents(1)
|
||
assert "a1" not in stale
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_by_tags(agent_repo):
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile", "test"], max_concurrent=2)
|
||
a3 = AgentInfo(agent_id="a3", endpoint="e3", agent_tags=["test"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.upsert(a2)
|
||
await agent_repo.upsert(a3)
|
||
|
||
ids = {a.agent_id for a in await agent_repo.by_tags(["compile"])}
|
||
assert ids == {"a1", "a2"}
|
||
|
||
ids = {a.agent_id for a in await agent_repo.by_tags(["compile", "test"])}
|
||
assert ids == {"a2"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_by_tags_excludes_offline_unavailable_stopping(agent_repo):
|
||
"""调度候选排除 offline/unavailable/stopping，允许 ready/processing（processing 仍有并发容量）。"""
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
|
||
a3 = AgentInfo(agent_id="a3", endpoint="e3", agent_tags=["compile"], max_concurrent=2)
|
||
a4 = AgentInfo(agent_id="a4", endpoint="e4", agent_tags=["compile"], max_concurrent=2)
|
||
a5 = AgentInfo(agent_id="a5", endpoint="e5", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.upsert(a2)
|
||
await agent_repo.upsert(a3)
|
||
await agent_repo.upsert(a4)
|
||
await agent_repo.upsert(a5)
|
||
await agent_repo.update_status("a2", AgentStatus.OFFLINE) # 离线（心跳失联）不可调度
|
||
await agent_repo.set_unavailable("a3") # 手动不可用不可调度
|
||
await agent_repo.update_status("a4", AgentStatus.STOPPING) # 停止中不可调度
|
||
await agent_repo.update_status("a5", AgentStatus.PROCESSING) # 处理中但还有容量，可调度
|
||
|
||
ids = {a.agent_id for a in await agent_repo.by_tags(["compile"])}
|
||
assert ids == {"a1", "a5"}
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_unavailable_heartbeat_not_revive(agent_repo):
|
||
"""手动置不可用（unavailable）后，即使发心跳也不复活。"""
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.set_unavailable("a1")
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
|
||
# 心跳不应复活
|
||
await agent_repo.beat("a1", 0)
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
|
||
# 恢复可用
|
||
await agent_repo.set_available("a1")
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.READY
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_unavailable_distinct_from_offline(agent_repo):
|
||
"""离线（心跳失联）与不可用（手动禁用）是两个独立状态。"""
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.upsert(a2)
|
||
# a1 心跳失联→离线；a2 手动禁用→不可用
|
||
await agent_repo.update_status("a1", AgentStatus.OFFLINE)
|
||
await agent_repo.set_unavailable("a2")
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE
|
||
assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE
|
||
# 发心跳：离线的 a1 恢复 ready，不可用的 a2 保持 unavailable
|
||
await agent_repo.beat("a1", 0)
|
||
await agent_repo.beat("a2", 0)
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.READY
|
||
assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mark_idle_keeps_unavailable(agent_repo):
|
||
"""任务结束后 mark_idle：unavailable 的 agent 保持不可用，其余置回 ready。"""
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.upsert(a2)
|
||
await agent_repo.set_unavailable("a1")
|
||
await agent_repo.update_status("a2", AgentStatus.PROCESSING)
|
||
await agent_repo.mark_idle("a1")
|
||
await agent_repo.mark_idle("a2")
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
|
||
assert (await agent_repo.get("a2")).status == AgentStatus.READY
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_heartbeat_recovers_stale_offline(agent_repo):
|
||
"""超时离线（非手动）的 agent 在发心跳后恢复为 ready。"""
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.update_status("a1", AgentStatus.OFFLINE)
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE
|
||
await agent_repo.beat("a1", 0)
|
||
assert (await agent_repo.get("a1")).status == AgentStatus.READY
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_heartbeat_keeps_busy_states(agent_repo):
|
||
"""processing/stopping 状态不因心跳改变。"""
|
||
for st in (AgentStatus.PROCESSING, AgentStatus.STOPPING):
|
||
a = AgentInfo(agent_id=f"a-{st.value}", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
|
||
await agent_repo.upsert(a)
|
||
await agent_repo.update_status(a.agent_id, st)
|
||
await agent_repo.beat(a.agent_id, 1)
|
||
assert (await agent_repo.get(a.agent_id)).status == st
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_remove(agent_repo):
|
||
a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"])
|
||
await agent_repo.upsert(a1)
|
||
await agent_repo.remove("a1")
|
||
assert await agent_repo.get("a1") is None
|
||
assert await agent_repo.by_tags(["compile"]) == [] |