import unittest
from modules.lawbot.rag_with_langchain import load_documents, split_documents
from langchain.schema import Document

class TestRAGWithLangchain(unittest.TestCase):

    def test_load_documents(self):
        """Test loading documents from the pickle file"""
        documents = load_documents()
        self.assertIsInstance(documents, list)
        self.assertGreater(len(documents), 0)
        self.assertIsInstance(documents[0], Document)
        self.assertIsNotNone(documents[0].page_content)

    def test_split_documents(self):
        """Test splitting documents into chunks"""
        test_documents = [
            Document(page_content="This is a test document with some content."),
            Document(page_content="Another document with different content for testing.")
        ]
        docs = split_documents(test_documents)
        self.assertIsInstance(docs, list)
        self.assertGreater(len(docs), 0)
        for doc in docs:
            self.assertIsInstance(doc, Document)
            self.assertLessEqual(len(doc.page_content), 1100)  # chunk_size + overlap

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