Testing

from tenxgraph.qa.testing import (
    InMemoryStore,
    MockMCPClient,
    MockToolRegistry,
    QuickTest,
    TestAgent,
    TestContext,
    TestResult,
)

from tenxgraph.qa.evaluation import QuickEval

from tenxgraph.qa.testing import TestAgent

agent = TestAgent(responses=["Hello from the test agent!", "How can I help?"])

# Run a one-node graph with a TestAgent instead of a real model.
import asyncio

from tenxgraph.core.graph import StateGraph
from tenxgraph.core.state import Message
from tenxgraph.qa.testing import TestAgent
from tenxgraph.utils import END


async def main() -> None:
    test_agent = TestAgent(responses=["Hello!"])

    graph = StateGraph()
    graph.add_node("MAIN", test_agent)
    graph.set_entry_point("MAIN")
    graph.add_edge("MAIN", END)

    app = graph.compile()
    result = await app.ainvoke({"messages": [Message.text_message("Hi")]})
    print(result["messages"][-1].text())
    test_agent.assert_called_times(1)


asyncio.run(main())

from tenxgraph.qa.testing import TestAgent


def get_weather(query: str) -> str:
    """Return a canned forecast."""
    return "sunny"


agent = TestAgent(responses=["It is sunny."], tools=[get_weather])

# Run a single turn and a tool scenario.
import asyncio

from tenxgraph.qa.testing import QuickTest


async def main() -> None:
    result = await QuickTest.single_turn(
        agent_response="The capital is Paris.",
        user_message="What is the capital of France?",
    )
    result.assert_contains("Paris")

    result = await QuickTest.with_tools(
        query="What's the weather in Paris?",
        response="It's sunny in Paris.",
        tools=["get_weather"],
        tool_responses={"get_weather": "sunny"},
    )
    result.assert_tool_called("get_weather")


asyncio.run(main())

result.assert_contains("Hello").assert_no_errors().assert_message_count(2)

from tenxgraph.core.graph import ToolNode
from tenxgraph.qa.testing import MockMCPClient

mock_mcp = MockMCPClient()
mock_mcp.add_tool(
    name="search_docs",
    description="Search the documentation",
    parameters={"query": {"type": "string", "description": "Search query"}},
    handler=lambda query: f"Found results for: {query}",
)

tool_node = ToolNode([], client=mock_mcp)

from tenxgraph.core.graph import ToolNode
from tenxgraph.qa.testing import MockToolRegistry

registry = MockToolRegistry()
registry.register("get_weather", lambda location: f"Sunny in {location}")

tools = ToolNode(registry.get_tool_list())

# ... run the graph ...

assert registry.was_called("get_weather")
assert registry.call_count("get_weather") == 1
assert registry.get_last_call("get_weather")["kwargs"]["location"] == "London"

from tenxgraph.core.graph import StateGraph
from tenxgraph.qa.testing import InMemoryStore

store = InMemoryStore()
# store.set_search_results([...]) makes every search return those results.

graph = StateGraph()
# ... add nodes ...
# app = graph.compile(store=store)

# Build and compile a graph inside an isolated test context.
from tenxgraph.qa.testing import TestContext
from tenxgraph.utils import END

with TestContext() as ctx:
    graph = ctx.create_graph()
    graph.add_node("MAIN", ctx.create_test_agent(responses=["Hello!"]))
    graph.set_entry_point("MAIN")
    graph.add_edge("MAIN", END)
    app = graph.compile(store=ctx.get_store())

# Quick checks against a graph compiled with a collector.
import asyncio

from tenxgraph.qa.evaluation import QuickEval
from tenxgraph.qa.evaluation.testing import create_eval_app


async def run(graph) -> None:
    app, collector = create_eval_app(graph)  # graph is an uncompiled StateGraph

    await QuickEval.check(
        graph=app,
        collector=collector,
        query="What is 2+2?",
        expected_response_contains="4",
    )

    await QuickEval.batch(
        graph=app,
        collector=collector,
        test_pairs=[("Hello", "Hi there"), ("What can you do?", "I can help")],
    )

    await QuickEval.tool_usage(
        graph=app,
        collector=collector,
        test_cases=[("What's the weather?", "It's sunny", ["get_weather"])],
    )