Spaces:
Sleeping
Sleeping
File size: 1,967 Bytes
f866820 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """
Unit tests for the RAG Document Assistant
"""
import unittest
from unittest.mock import patch, MagicMock
class TestOrchestrator(unittest.TestCase):
@patch('src.orchestrator.pinecone_search')
@patch('src.orchestrator.call_llm')
def test_orchestrate_query_success(self, mock_call_llm, mock_pinecone_search):
"""Test successful query orchestration"""
# Mock the retrieval function
mock_pinecone_search.return_value = [
{
"id": "doc1",
"score": 0.85,
"text": "This is a sample document chunk about GDPR regulations.",
"metadata": {"source": "gdpr.pdf"}
}
]
# Mock the LLM function
mock_call_llm.return_value = {
"text": "GDPR is a regulation about data protection.",
"meta": {"provider": "gemini", "model": "gemini-1.5-flash"}
}
# Import after mocking
from src.orchestrator import orchestrate_query
# Test the function
result = orchestrate_query("What is GDPR?", top_k=1)
# Assertions
self.assertIn('answer', result)
self.assertIn('sources', result)
self.assertIn('citations', result)
self.assertEqual(result['answer'], "GDPR is a regulation about data protection.")
# Verify mocks were called
mock_pinecone_search.assert_called_once()
mock_call_llm.assert_called_once()
def test_orchestrate_query_invalid_input(self):
"""Test query orchestration with invalid input"""
from src.orchestrator import orchestrate_query
# Test with empty query
result = orchestrate_query("")
self.assertEqual(result['answer'], "")
# Test with None query
result = orchestrate_query(None)
self.assertEqual(result['answer'], "")
if __name__ == '__main__':
unittest.main() |