Function test_langgraph_agent

tests/unittests/agents/test_langgraph_agent.py:219–264

(
    checkpointer_value, events_list, expected_messages
)

)
@pytest.mark.asyncio
async def test_langgraph_agent(
    checkpointer_value, events_list, expected_messages
):
  mock_graph = MagicMock(spec=CompiledGraph)
  mock_graph_state = MagicMock()
  mock_graph_state.values = {}
  mock_graph.get_state.return_value = mock_graph_state

  mock_graph.checkpointer = checkpointer_value
  mock_graph.invoke.return_value = {
      "messages": [AIMessage(content="test response")]
  }

  mock_parent_context = MagicMock(spec=InvocationContext)
  mock_parent_context._state_schema = None
  mock_session = MagicMock()
  mock_parent_context.session = mock_session
  mock_parent_context.user_content = types.Content(
      role="user", parts=[types.Part.from_text(text="test prompt")]
  )
  mock_parent_context.branch = "parent_agent"
  mock_parent_context.end_invocation = False
  mock_session.events = events_list
  mock_parent_context.invocation_id = "test_invocation_id"
  mock_parent_context.model_copy.return_value = mock_parent_context
  mock_parent_context.plugin_manager = PluginManager(plugins=[])

  weather_agent = LangGraphAgent(
      name="weather_agent",
      description="A agent that answers weather questions",
      instruction="test system prompt",
      graph=mock_graph,
  )

  result_event = None
  async for event in weather_agent.run_async(mock_parent_context):
    result_event = event

  assert result_event.author == "weather_agent"
  assert result_event.content.parts[0].text == "test response"

  mock_graph.invoke.assert_called_once()
  mock_graph.invoke.assert_called_with(
      {"messages": expected_messages},
      {"configurable": {"thread_id": mock_session.id}},
  )