Repository: emcie-co/parlant
File: server/tests/core/persistence/test_json_file_storage.py (Python)

Code before:
    container: Container
    agent_id: AgentId
    sync_await: SyncAwaiter


@fixture
def context(
    container: Container,
    agent_id: AgentId,
    sync_await: SyncAwaiter,
) -> _TestContext:
    return _TestContext(container, agent_id, sync_await)


@fixture
async def new_file() -> AsyncIterator[Path]:
    with tempfile.NamedTemporaryFile() as file:
        yield Path(file.name)


async def test_agent_creation(
    new_file: Path,
) -> None:
    async with JSONFileDocumentDatabase(new_file) as agent_db:
        agent_store = AgentDocumentStore(agent_db)
        agent_without_description = await agent_store.create_agent(name="Test Agent")
        agent_with_description = await agent_store.create_agent(
            name="Test Agent",
            description="You are a test agent",
        )

        agents = list(await agent_store.list_agents())

        assert len(agents) == 2
        assert agents[0] == agent_without_description

    with open(new_file) as f:
        agents_from_json = json.load(f)

    assert len(agents_from_json["agents"]) == 2

    json_agent_without_description = agents_from_json["agents"][0]
    assert json_agent_without_description["id"] == agent_without_description.id
    assert json_agent_without_description["name"] == agent_without_description.name
    assert json_agent_without_description["description"] is None
    assert (
        datetime.fromisoformat(json_agent_without_description["creation_utc"])
        == agent_without_description.creation_utc
    )

    json_agent_with_description = agents_from_json["agents"][1]

Review comment (suggestion):
Parameterize this test on two agent specs, instead of copy/pasting into one test.

This is important partly because, in principle, we probably shouldn't support two agents with the same name anyway. But the test will be cleaner anyhow this way.

Code after:
    container: Container
    agent_id: AgentId
    sync_await: SyncAwaiter


@fixture
def context(
    container: Container,
    agent_id: AgentId,
    sync_await: SyncAwaiter,
) -> _TestContext:
    return _TestContext(container, agent_id, sync_await)


@fixture
async def new_file() -> AsyncIterator[Path]:
    with tempfile.NamedTemporaryFile() as file:
        yield Path(file.name)


@mark.parametrize(
    ("agent_configuration"),
    [
        ({"name": "Test Agent"}),
        ({"name": "Test Agent", "description": "You are a test agent"}),
    ],
)
async def test_agent_creation(
    new_file: Path,
    agent_configuration: dict[str, Any],
) -> None:
    async with JSONFileDocumentDatabase(new_file) as agent_db:
        agent_store = AgentDocumentStore(agent_db)
        agent = await agent_store.create_agent(**agent_configuration)

        agents = list(await agent_store.list_agents())

        assert len(agents) == 1
        assert agents[0] == agent

    with open(new_file) as f:
        agents_from_json = json.load(f)

    assert len(agents_from_json["agents"]) == 1

    json_agent = agents_from_json["agents"][0]
    assert json_agent["id"] == agent.id
    assert json_agent["name"] == agent.name
    assert json_agent["description"] == agent.description
    assert datetime.fromisoformat(json_agent["creation_utc"]) == agent.creation_utc