im-amrith/Rag-engine-backend

1import sys2from unittest.mock import MagicMock3import os4import json5import unittest6 7# Mock sentence_transformers before importing rag_engine8sys.modules["sentence_transformers"] = MagicMock()9 10# Also mock psycopg2 if needed, but we want to test DB interaction if possible.11# However, if we want to avoid DB connection issues, we can mock psycopg2 too.12# But testing the SQL query is the main point.13# Let's assume psycopg2 is available (we installed it).14# We need to make sure RAGEngine doesn't fail on __init__ due to model load.15# The mocked SentenceTransformer should handle it.16 17# We need to set DATABASE_URL18os.environ["DATABASE_URL"] = "postgresql://postgres:password@localhost:5432/rag_engine" # Dummy or real?19# If we use real DB, we need real connection.20# If we mock DB, we verify the SQL string.21 22# Let's try to use the real DB if available, otherwise mock.23# The user has a DB_URL in .env?24from dotenv import load_dotenv25load_dotenv()26 27# If we can't connect to DB, we mock psycopg228try:29    import psycopg230    conn = psycopg2.connect(os.getenv("DATABASE_URL"))31    conn.close()32    USE_REAL_DB = True33except:34    USE_REAL_DB = False35    print("Could not connect to real DB, using mocks.")36 37# Now import38sys.path.append(os.getcwd())39try:40    from rag_engine import RAGEngine41except ImportError:42    # If it fails due to other imports43    print("Failed to import RAGEngine")44    sys.exit(1)45 46class TestRAGEngine(unittest.TestCase):47    def setUp(self):48        if USE_REAL_DB:49            self.engine = RAGEngine()50            # We need to mock the model attribute since we mocked the class51            self.engine.model = MagicMock()52            self.engine.model.encode.return_value.tolist.return_value = [0.1] * 38453            54            # Create a test user55            self.test_email = "unittest@example.com"56            self.user_id = self.engine.create_user(self.test_email, "hash")57            if not self.user_id:58                # User might exist59                user = self.engine.get_user(self.test_email)60                self.user_id = user[0]61        else:62            self.engine = RAGEngine()63            self.engine.conn = MagicMock()64            self.user_id = 165 66    def test_get_chat_history_and_item(self):67        # 1. Save a chat68        self.engine.save_chat("Unit Test User", "Unit Test AI", self.user_id)69        70        # 2. Get History71        history = self.engine.get_chat_history(self.user_id, limit=1)72        self.assertTrue(len(history) > 0)73        item = history[0]74        self.assertIn("id", item)75        self.assertEqual(item["user"], "Unit Test User")76        77        chat_id = item["id"]78        79        # 3. Get Specific Item80        fetched_item = self.engine.get_chat_item(chat_id, self.user_id)81        self.assertIsNotNone(fetched_item)82        self.assertEqual(fetched_item["id"], chat_id)83        self.assertEqual(fetched_item["user"], "Unit Test User")84        self.assertEqual(fetched_item["ai"], "Unit Test AI")85        86        print("Test passed!")87 88if __name__ == "__main__":89    unittest.main()90