You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
73 lines
2.5 KiB
73 lines
2.5 KiB
"""
|
|
Unit tests for RAG pipeline
|
|
"""
|
|
import pytest
|
|
from unittest.mock import Mock, patch
|
|
from src.knowledge.rag_pipeline import create_knowledge_base, ingest_excel_data
|
|
|
|
|
|
class TestRAGPipeline:
|
|
"""Test cases for RAG pipeline"""
|
|
|
|
@patch('src.knowledge.rag_pipeline.get_qdrant_store')
|
|
def test_create_knowledge_base_qdrant(self, mock_get_qdrant):
|
|
"""Test creating knowledge base with Qdrant"""
|
|
mock_vector_db = Mock()
|
|
mock_get_qdrant.return_value = mock_vector_db
|
|
|
|
kb = create_knowledge_base(vector_store_type="qdrant")
|
|
|
|
mock_get_qdrant.assert_called_once()
|
|
assert kb.vector_db == mock_vector_db
|
|
|
|
@patch('src.knowledge.rag_pipeline.get_pgvector_store')
|
|
def test_create_knowledge_base_pgvector(self, mock_get_pgvector):
|
|
"""Test creating knowledge base with PgVector"""
|
|
mock_vector_db = Mock()
|
|
mock_get_pgvector.return_value = mock_vector_db
|
|
|
|
kb = create_knowledge_base(vector_store_type="pgvector")
|
|
|
|
mock_get_pgvector.assert_called_once()
|
|
assert kb.vector_db == mock_vector_db
|
|
|
|
def test_create_knowledge_base_invalid_type(self):
|
|
"""Test creating knowledge base with invalid type"""
|
|
with pytest.raises(ValueError, match="Unsupported vector store type"):
|
|
create_knowledge_base(vector_store_type="invalid")
|
|
|
|
@patch('pandas.read_excel')
|
|
def test_ingest_hadiths_data(self, mock_read_excel, sample_hadith_data):
|
|
"""Test ingesting hadith data"""
|
|
# Mock DataFrame
|
|
mock_df = Mock()
|
|
mock_df.iterrows.return_value = [(0, sample_hadith_data)]
|
|
mock_read_excel.return_value = mock_df
|
|
|
|
# Mock knowledge base
|
|
mock_kb = Mock()
|
|
|
|
# Mock file operations
|
|
with patch('builtins.open', Mock()):
|
|
count = ingest_excel_data(mock_kb, "test.xlsx", "hadiths")
|
|
|
|
assert count == 1
|
|
mock_kb.add_content.assert_called_once()
|
|
|
|
@patch('pandas.read_excel')
|
|
def test_ingest_articles_data(self, mock_read_excel, sample_article_data):
|
|
"""Test ingesting article data"""
|
|
# Mock DataFrame
|
|
mock_df = Mock()
|
|
mock_df.iterrows.return_value = [(0, sample_article_data)]
|
|
mock_read_excel.return_value = mock_df
|
|
|
|
# Mock knowledge base
|
|
mock_kb = Mock()
|
|
|
|
# Mock file operations
|
|
with patch('builtins.open', Mock()):
|
|
count = ingest_excel_data(mock_kb, "test.xlsx", "articles")
|
|
|
|
assert count == 1
|
|
mock_kb.add_content.assert_called_once()
|