From a404b73ccc115f4101fb1419eb0f6e757153672d Mon Sep 17 00:00:00 2001
From: AquibPy <mohdaquib31@gmail.com>
Date: Wed, 15 May 2024 16:07:33 +0530
Subject: [PATCH] ADDED: test case for Advance RAG System endpoint

---
 test_main.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/test_main.py b/test_main.py
index b7a50ec..4c3aa21 100644
--- a/test_main.py
+++ b/test_main.py
@@ -7,7 +7,6 @@
 client = TestClient(app)
 # db = MongoDB()
 
-
 def test_invoice_extractor():
     with patch('api.MongoDB') as MockMongoDB:
         mock_insert_data = MagicMock(return_value={'Status': 'Successfully Inserted!!!', 'Document_id': 'mock_id'})
@@ -177,4 +176,16 @@ def test_qa_url_doc_with_document():
             assert response.status_code == 200
             response_data = response.json()
             assert "response" in response_data
-            assert isinstance(response_data["response"], str)
\ No newline at end of file
+            assert isinstance(response_data["response"], str)
+
+def test_advance_rag():
+    pdf_file = open("data/yolo.pdf", "rb")
+    data = {"question": "Summary in 200 words" ,"model": "llama3-70b-8192"}
+
+    with patch('api.MongoDB') as MockMongoDB:
+        mock_insert_data = MagicMock(return_value={'Status': 'Successfully Inserted!!!', 'Document_id': 'mock_id'})
+        MockMongoDB.return_value.insert_data = mock_insert_data
+        
+        response = client.post("/advance_rag_llama_index", files={"pdf": pdf_file}, data=data)
+        assert response.status_code == 200
+        assert "response" in response.json()
\ No newline at end of file