-
Notifications
You must be signed in to change notification settings - Fork 16
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Marc Fabian Mezger
committed
Jul 1, 2024
1 parent
acbe98c
commit 5d1f5e2
Showing
3 changed files
with
97 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -130,6 +130,7 @@ exclude = [ | |
"node_modules", | ||
"site-packages", | ||
"venv", | ||
"tests", | ||
] | ||
|
||
# Same as Black. | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
from unittest.mock import patch | ||
|
||
import pytest | ||
from cohere_service import CohereService | ||
from omegaconf import DictConfig | ||
|
||
from agent.data_model.request_data_model import SearchParams | ||
|
||
|
||
@pytest.fixture() | ||
def mock_cfg(): | ||
return DictConfig({"qdrant": {"collection_name_cohere": "test_collection"}, "cohere_embeddings": {"embedding_model_name": "test-model", "size": 768}}) | ||
|
||
|
||
@pytest.fixture() | ||
def cohere_service_instance(mock_cfg): | ||
with patch("cohere_service.load_dotenv"), patch("cohere_service.load_config", return_value=mock_cfg): | ||
return CohereService(cfg=mock_cfg, collection_name=None) | ||
|
||
|
||
def test_cohere_service_initialization(cohere_service_instance) -> None: | ||
assert cohere_service_instance.collection_name == "test_collection" | ||
assert cohere_service_instance.cfg.cohere_embeddings.embedding_model_name == "test-model" | ||
|
||
|
||
@patch("cohere_service.DirectoryLoader") | ||
@patch("cohere_service.PyPDFium2Loader") | ||
@patch("cohere_service.TextLoader") | ||
def test_embed_documents(mock_text_loader, mock_pdf_loader, mock_dir_loader, cohere_service_instance) -> None: | ||
cohere_service_instance.embed_documents(directory="tests/resources/", file_ending=".pdf") | ||
mock_dir_loader.assert_called_once() | ||
# Further assertions can be added to verify the behavior | ||
|
||
|
||
def test_create_collection(cohere_service_instance) -> None: | ||
with patch("cohere_service.generate_collection") as mock_generate_collection: | ||
result = cohere_service_instance.create_collection(name="new_collection") | ||
mock_generate_collection.assert_called_once_with("new_collection", 768) | ||
assert result is True | ||
|
||
|
||
def test_create_search_chain(cohere_service_instance) -> None: | ||
search_params = SearchParams(query="test query", k=5) | ||
with patch("cohere_service.chain") as mock_chain: | ||
cohere_service_instance.create_search_chain(search=search_params) | ||
# Verify the chain decorator was used | ||
mock_chain.assert_called_once() | ||
# Further assertions can be added to verify the retriever's behavior | ||
|
||
|
||
# Additional tests can be added for other methods and edge cases |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,45 @@ | ||
from unittest.mock import MagicMock, patch | ||
|
||
import pytest | ||
from ollama_service import OllamaService | ||
from omegaconf import DictConfig | ||
|
||
|
||
@pytest.fixture() | ||
def mock_cfg(): | ||
return DictConfig({"qdrant": {"collection_name_ollama": "test_collection", "embeddings_size": 768}, "ollama_embeddings": {"embedding_model_name": "test_model"}}) | ||
|
||
|
||
@pytest.fixture() | ||
def ollama_service(mock_cfg): | ||
with patch("ollama_service.load_config", return_value=mock_cfg): | ||
return OllamaService(cfg=mock_cfg, collection_name="test_collection") | ||
|
||
|
||
def test_initialization(ollama_service) -> None: | ||
assert ollama_service.collection_name == "test_collection" | ||
assert ollama_service.cfg == mock_cfg() | ||
|
||
|
||
@patch("ollama_service.DirectoryLoader") | ||
@patch("ollama_service.PyPDFium2Loader") | ||
@patch("ollama_service.TextLoader") | ||
def test_embed_documents(mock_text_loader, mock_pdf_loader, mock_dir_loader, ollama_service) -> None: | ||
ollama_service.embed_documents(directory="tests/resources/", file_ending=".pdf") | ||
mock_dir_loader.assert_called_once() | ||
# Further assertions can be added to validate the behavior | ||
|
||
|
||
def test_create_collection(ollama_service) -> None: | ||
with patch.object(ollama_service.vector_db, "add_texts", return_value=True): | ||
result = ollama_service.create_collection(name="new_collection") | ||
assert result is True | ||
# You can add more assertions here to check if the collection was created with the correct parameters | ||
|
||
|
||
@patch("ollama_service.chain") | ||
def test_create_search_chain(mock_chain, ollama_service) -> None: | ||
search_params = MagicMock() | ||
ollama_service.create_search_chain(search=search_params) | ||
mock_chain.assert_called_once() | ||
# Further assertions can be added to validate the search chain behavior |