agency-swarm-cn-git/tests/test_agent_modules/test_agent_flow.py

from __future__ import annotations

import operator

import pytest

from agency_swarm import Agent
from agency_swarm.agent.agent_flow import AgentFlow


def _agent(name: str) -> Agent:
    return Agent(name=name, instructions=f"Instructions for {name}", model="gpt-5.4-mini")


def test_agent_flow_gt_rejects_non_agent() -> None:
    flow = AgentFlow([_agent("A"), _agent("B")])
    with pytest.raises(TypeError, match="Can only chain to Agent instances"):
        operator.gt(flow, object())


def test_agent_flow_lt_rejects_non_agent() -> None:
    flow = AgentFlow([_agent("A"), _agent("B")])
    with pytest.raises(TypeError, match="Can only chain to Agent instances"):
        operator.lt(flow, object())


def test_agent_flow_chain_and_repr() -> None:
    agent_a = _agent("A")
    agent_b = _agent("B")
    agent_c = _agent("C")
    agent_d = _agent("D")

    flow = AgentFlow([agent_a, agent_b])
    expanded = flow > agent_c
    prepended = expanded < agent_d

    assert expanded.get_all_flows() == [(agent_a, agent_b), (agent_b, agent_c)]
    assert prepended.get_all_flows() == [(agent_d, agent_a), (agent_a, agent_b), (agent_b, agent_c)]
    assert repr(expanded) == "AgentFlow(A > B > C)"


def test_agent_flow_bool_tracks_and_clears_chain_flows() -> None:
    AgentFlow.get_and_clear_chain_flows()
    agent_a = _agent("A")
    agent_b = _agent("B")
    agent_c = _agent("C")

    flow = AgentFlow([agent_a, agent_b, agent_c])

    assert bool(flow) is True
    assert bool(flow) is True  # duplicate truthiness should not duplicate tracked flows

    tracked = AgentFlow.get_and_clear_chain_flows()
    assert tracked == [(agent_a, agent_b), (agent_b, agent_c)]
    assert AgentFlow.get_and_clear_chain_flows() == []