moss/examples/cookbook/langchain/test_integration.py at main · usemoss/moss · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import unittest
from unittest.mock import AsyncMock, patch, MagicMock
from moss_langchain import MossRetriever, get_moss_tool
from langchain_core.documents import Document


class TestMossRetriever(unittest.IsolatedAsyncioTestCase):
    """Behavioral tests for MossRetriever."""

    @patch('moss_langchain.MossClient')
    async def test_aget_relevant_documents_maps_results(self, mock_client_cls):
        mock_client = mock_client_cls.return_value
        mock_client.load_index = AsyncMock()
        mock_client.query = AsyncMock(return_value=MagicMock(
            docs=[
                MagicMock(text="hello world", score=0.9, id="d1"),
                MagicMock(text="foo bar", score=0.7, id="d2"),
            ]
        ))

        retriever = MossRetriever(
            project_id="p", project_key="k", index_name="idx"
        )
        docs = await retriever._aget_relevant_documents("test query")

        self.assertEqual(len(docs), 2)
        self.assertIsInstance(docs[0], Document)
        self.assertEqual(docs[0].page_content, "hello world")
        self.assertEqual(docs[0].metadata, {"score": 0.9, "id": "d1"})
        self.assertEqual(docs[1].page_content, "foo bar")

    @patch('moss_langchain.MossClient')
    async def test_aget_relevant_documents_empty_results(self, mock_client_cls):
        mock_client = mock_client_cls.return_value
        mock_client.load_index = AsyncMock()
        mock_client.query = AsyncMock(return_value=MagicMock(docs=[]))

        retriever = MossRetriever(
            project_id="p", project_key="k", index_name="idx"
        )
        docs = await retriever._aget_relevant_documents("empty query")

        self.assertEqual(docs, [])

    @patch('moss_langchain.MossClient')
    async def test_load_index_called_once(self, mock_client_cls):
        mock_client = mock_client_cls.return_value
        mock_client.load_index = AsyncMock()
        mock_client.query = AsyncMock(return_value=MagicMock(docs=[]))

        retriever = MossRetriever(
            project_id="p", project_key="k", index_name="idx"
        )

        await retriever._aget_relevant_documents("q1")
        await retriever._aget_relevant_documents("q2")

        mock_client.load_index.assert_awaited_once_with("idx")


class TestGetMossTool(unittest.IsolatedAsyncioTestCase):
    """Behavioral tests for get_moss_tool."""

    @patch('moss_langchain.MossClient')
    async def test_coroutine_formats_output(self, mock_client_cls):
        mock_client = mock_client_cls.return_value
        mock_client.load_index = AsyncMock()
        mock_client.query = AsyncMock(return_value=MagicMock(
            docs=[
                MagicMock(text="first result", score=0.9, id="d1"),
                MagicMock(text="second result", score=0.8, id="d2"),
            ]
        ))

        tool = get_moss_tool("p", "k", "idx")
        result = await tool.coroutine("search query")

        self.assertIn("Result 1:\nfirst result", result)
        self.assertIn("Result 2:\nsecond result", result)

    @patch('moss_langchain.MossClient')
    async def test_coroutine_empty_results(self, mock_client_cls):
        mock_client = mock_client_cls.return_value
        mock_client.load_index = AsyncMock()
        mock_client.query = AsyncMock(return_value=MagicMock(docs=[]))

        tool = get_moss_tool("p", "k", "idx")
        result = await tool.coroutine("no results query")

        self.assertEqual(result, "No relevant information found.")


if __name__ == '__main__':
    unittest.main()