fleet-memory/hindsight-integrations/langgraph/tests/test_nodes.py

"""Unit tests for Hindsight LangGraph nodes."""

from unittest.mock import AsyncMock, MagicMock

import pytest
from hindsight_langgraph import create_recall_node, create_retain_node
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage


def _mock_client():
    client = MagicMock()
    client.aretain = AsyncMock()
    client.arecall = AsyncMock()
    return client


def _mock_recall_response(texts: list[str]):
    response = MagicMock()
    results = []
    for t in texts:
        r = MagicMock()
        r.text = t
        results.append(r)
    response.results = results
    return response


class TestRecallNode:
    @pytest.mark.asyncio
    async def test_injects_memories_as_system_message(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response(
            ["User likes Python", "User is in NYC"]
        )
        node = create_recall_node(bank_id="test-bank", client=client)

        state = {"messages": [HumanMessage(content="What do you remember about me?")]}
        result = await node(state)

        assert len(result["messages"]) == 1
        msg = result["messages"][0]
        assert isinstance(msg, SystemMessage)
        assert "User likes Python" in msg.content
        assert "User is in NYC" in msg.content

    @pytest.mark.asyncio
    async def test_returns_empty_when_no_human_message(self):
        client = _mock_client()
        node = create_recall_node(bank_id="test-bank", client=client)

        state = {"messages": [SystemMessage(content="You are a helpful assistant")]}
        result = await node(state)

        assert result["messages"] == []
        client.arecall.assert_not_called()

    @pytest.mark.asyncio
    async def test_returns_empty_when_no_results(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response([])
        node = create_recall_node(bank_id="test-bank", client=client)

        state = {"messages": [HumanMessage(content="hello")]}
        result = await node(state)

        assert result["messages"] == []

    @pytest.mark.asyncio
    async def test_respects_max_results(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response(
            ["fact1", "fact2", "fact3", "fact4", "fact5"]
        )
        node = create_recall_node(bank_id="test-bank", client=client, max_results=2)

        state = {"messages": [HumanMessage(content="query")]}
        result = await node(state)

        msg = result["messages"][0]
        assert "1. fact1" in msg.content
        assert "2. fact2" in msg.content
        assert "3." not in msg.content

    @pytest.mark.asyncio
    async def test_resolves_bank_id_from_config(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response(["fact"])
        node = create_recall_node(client=client, bank_id_from_config="user_id")

        state = {"messages": [HumanMessage(content="hello")]}
        config = {"configurable": {"user_id": "user-456"}}
        await node(state, config=config)

        call_kwargs = client.arecall.call_args[1]
        assert call_kwargs["bank_id"] == "user-456"

    @pytest.mark.asyncio
    async def test_skips_when_no_bank_id(self):
        client = _mock_client()
        node = create_recall_node(client=client)

        state = {"messages": [HumanMessage(content="hello")]}
        result = await node(state)

        assert result["messages"] == []
        client.arecall.assert_not_called()

    @pytest.mark.asyncio
    async def test_handles_recall_error_gracefully(self):
        client = _mock_client()
        client.arecall.side_effect = RuntimeError("connection refused")
        node = create_recall_node(bank_id="test-bank", client=client)

        state = {"messages": [HumanMessage(content="hello")]}
        result = await node(state)

        assert result["messages"] == []

    @pytest.mark.asyncio
    async def test_passes_tags(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response(["fact"])
        node = create_recall_node(
            bank_id="test-bank",
            client=client,
            tags=["scope:user"],
            tags_match="all",
        )

        state = {"messages": [HumanMessage(content="hello")]}
        await node(state)

        call_kwargs = client.arecall.call_args[1]
        assert call_kwargs["tags"] == ["scope:user"]
        assert call_kwargs["tags_match"] == "all"


class TestRecallNodeOutputKey:
    @pytest.mark.asyncio
    async def test_output_key_returns_memory_text(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response(
            ["User likes Python", "User is in NYC"]
        )
        node = create_recall_node(
            bank_id="test-bank", client=client, output_key="memory_context"
        )

        state = {"messages": [HumanMessage(content="What do you remember?")]}
        result = await node(state)

        assert "messages" not in result
        assert "memory_context" in result
        assert "User likes Python" in result["memory_context"]
        assert "User is in NYC" in result["memory_context"]

    @pytest.mark.asyncio
    async def test_output_key_returns_none_when_no_results(self):
        client = _mock_client()
        client.arecall.return_value = _mock_recall_response([])
        node = create_recall_node(
            bank_id="test-bank", client=client, output_key="memory_context"
        )

        state = {"messages": [HumanMessage(content="hello")]}
        result = await node(state)

        assert result == {"memory_context": None}

    @pytest.mark.asyncio
    async def test_output_key_returns_none_on_error(self):
        client = _mock_client()
        client.arecall.side_effect = RuntimeError("connection refused")
        node = create_recall_node(
            bank_id="test-bank", client=client, output_key="memory_context"
        )

        state = {"messages": [HumanMessage(content="hello")]}
        result = await node(state)

        assert result == {"memory_context": None}


class TestRetainNode:
    @pytest.mark.asyncio
    async def test_retains_human_messages(self):
        client = _mock_client()
        node = create_retain_node(bank_id="test-bank", client=client)

        state = {
            "messages": [
                HumanMessage(content="I like pizza"),
                AIMessage(content="Got it!"),
            ]
        }
        await node(state)

        client.aretain.assert_called_once()
        call_kwargs = client.aretain.call_args[1]
        assert call_kwargs["bank_id"] == "test-bank"
        assert "I like pizza" in call_kwargs["content"]
        assert "Got it!" not in call_kwargs["content"]

    @pytest.mark.asyncio
    async def test_retains_both_when_configured(self):
        client = _mock_client()
        node = create_retain_node(
            bank_id="test-bank", client=client, retain_human=True, retain_ai=True
        )

        state = {
            "messages": [
                HumanMessage(content="I like pizza"),
                AIMessage(content="Got it!"),
            ]
        }
        await node(state)

        call_kwargs = client.aretain.call_args[1]
        assert "I like pizza" in call_kwargs["content"]
        assert "Got it!" in call_kwargs["content"]

    @pytest.mark.asyncio
    async def test_skips_when_no_messages_match(self):
        client = _mock_client()
        node = create_retain_node(
            bank_id="test-bank", client=client, retain_human=False, retain_ai=False
        )

        state = {"messages": [HumanMessage(content="hello")]}
        await node(state)

        client.aretain.assert_not_called()

    @pytest.mark.asyncio
    async def test_passes_tags(self):
        client = _mock_client()
        node = create_retain_node(
            bank_id="test-bank", client=client, tags=["source:chat"]
        )

        state = {"messages": [HumanMessage(content="hello")]}
        await node(state)

        call_kwargs = client.aretain.call_args[1]
        assert call_kwargs["tags"] == ["source:chat"]

    @pytest.mark.asyncio
    async def test_resolves_bank_id_from_config(self):
        client = _mock_client()
        node = create_retain_node(client=client, bank_id_from_config="user_id")

        state = {"messages": [HumanMessage(content="hello")]}
        config = {"configurable": {"user_id": "user-789"}}
        await node(state, config=config)

        call_kwargs = client.aretain.call_args[1]
        assert call_kwargs["bank_id"] == "user-789"

    @pytest.mark.asyncio
    async def test_handles_retain_error_gracefully(self):
        client = _mock_client()
        client.aretain.side_effect = RuntimeError("connection refused")
        node = create_retain_node(bank_id="test-bank", client=client)

        state = {"messages": [HumanMessage(content="hello")]}
        # Should not raise
        result = await node(state)
        assert result["messages"] == []