"""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"]) == []