fix: typing error while test

tests/test_retriever.py

@@ -2,26 +2,24 @@
 from unittest.mock import patch, MagicMock
 from langrade import create_retriever
 from conftest import get_api_key
+from langchain.schema import Document
 
 
 class TestRetriever(unittest.TestCase):
     def setUp(self):
         self.api_key = get_api_key()
         self.urls = [
             "https://lilianweng.github.io/posts/2023-06-23-agent/",
-            "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",  # noqa: E501
+            "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
             "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
         ]
 
     @patch("langrade.retriever.WebBaseLoader")
     @patch("langrade.retriever.Chroma")
     def test_create_retriever(self, mock_chroma, mock_web_loader):
-        mock_doc = MagicMock()
-        mock_doc.page_content = "This is a test document"
+        mock_doc = Document(page_content="This is a test document", metadata={})
         mock_web_loader.return_value.load.return_value = [mock_doc]
-        mock_chroma.from_documents.return_value.as_retriever.return_value = (
-            MagicMock()
-        )  # noqa: E501
+        mock_chroma.from_documents.return_value.as_retriever.return_value = MagicMock()
 
         retriever = create_retriever(self.urls, self.api_key)
 
@@ -31,11 +29,8 @@ def test_create_retriever(self, mock_chroma, mock_web_loader):
 
     @patch("langrade.retriever.WebBaseLoader")
     @patch("langrade.retriever.Chroma")
-    def test_retriever_get_relevant_documents(
-        self, mock_chroma, mock_web_loader
-    ):  # noqa: E501
-        mock_doc = MagicMock()
-        mock_doc.page_content = "This is a test document"
+    def test_retriever_get_relevant_documents(self, mock_chroma, mock_web_loader):
+        mock_doc = Document(page_content="This is a test document", metadata={})
         mock_web_loader.return_value.load.return_value = [mock_doc]
         mock_retriever = MagicMock()
         mock_chroma.from_documents.return_value.as_retriever.return_value = (
@@ -45,7 +40,7 @@ def test_retriever_get_relevant_documents(
         retriever = create_retriever(self.urls, self.api_key)
         question = "What is AI?"
         mock_retriever.get_relevant_documents.return_value = [
-            MagicMock(page_content="AI is a field of computer science.")
+            Document(page_content="AI is a field of computer science.", metadata={})
         ]
 
         docs = retriever.get_relevant_documents(question)