55 lines
1.8 KiB
Python
55 lines
1.8 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.ONLINE
|
|
|
|
|
|
@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_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"]) == [] |