Mocking nodes with pytest

[baswenneker]

I can't get this working: mocking a node in my graph.

I'd like to patch get_messages_node in agents.core.generate_title_agent.
In the Pytest:

from unittest.mock import patch
import pytest
from agents.core.generate_title_agent import ainvoke, get_messages_node, AgentState, AIMessage

@pytest.mark.asyncio
async def test_mock_get_messages_node():
    with patch("agents.core.generate_title_agent.get_messages_node") as mock_node:
        mock_node.return_value = AgentState({"messages": [AIMessage(content="Mocked content")]})
        state = AgentState(thread_id="test_thread_123")
        result = await ainvoke(state)
        mock_node.assert_called_once()
        assert result["info"] == "Mocked content"

But mock_node.assert_called_once() fails.

Any ideas how to overcome this?