From b669a7df80f2d0aba5e83cdfae02a8b4c5e46d13 Mon Sep 17 00:00:00 2001 From: seehi <6580@pm.me> Date: Thu, 7 Mar 2024 21:41:28 +0800 Subject: [PATCH] unittest --- tests/metagpt/rag/engines/test_simple.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/metagpt/rag/engines/test_simple.py b/tests/metagpt/rag/engines/test_simple.py index fd34b4e8d..4125d480a 100644 --- a/tests/metagpt/rag/engines/test_simple.py +++ b/tests/metagpt/rag/engines/test_simple.py @@ -1,6 +1,6 @@ import pytest from llama_index.core import VectorStoreIndex -from llama_index.core.schema import NodeWithScore, TextNode +from llama_index.core.schema import Document, TextNode from metagpt.rag.engines import SimpleEngine from metagpt.rag.retrievers.base import ModifiableRAGRetriever @@ -37,7 +37,10 @@ class TestSimpleEngine: mock_get_response_synthesizer, ): # Mock - mock_simple_directory_reader.return_value.load_data.return_value = ["document1", "document2"] + mock_simple_directory_reader.return_value.load_data.return_value = [ + Document(text="document1"), + Document(text="document2"), + ] mock_get_retriever.return_value = mocker.MagicMock() mock_get_rankers.return_value = [mocker.MagicMock()] mock_get_response_synthesizer.return_value = mocker.MagicMock() @@ -114,7 +117,10 @@ class TestSimpleEngine: def test_add_docs(self, mocker): # Mock mock_simple_directory_reader = mocker.patch("metagpt.rag.engines.simple.SimpleDirectoryReader") - mock_simple_directory_reader.return_value.load_data.return_value = ["document1", "document2"] + mock_simple_directory_reader.return_value.load_data.return_value = [ + Document(text="document1"), + Document(text="document2"), + ] mock_retriever = mocker.MagicMock(spec=ModifiableRAGRetriever)