Spaces:
Paused
Paused
| import pytest | |
| from unittest.mock import MagicMock, patch | |
| from helpers import create_prompt, find_similar_documents, load_vector_store, MANUAL_PATH | |
| def test_load_vector_store(): | |
| with patch("helpers.Chroma") as mocked_db: | |
| mocked_db.return_value = MagicMock() | |
| db = load_vector_store(MANUAL_PATH) | |
| assert db is not None | |
| mocked_db.from_documents.assert_called_once() | |
| def test_find_similar_documents(): | |
| # 1. Setup: Create a mock vector store | |
| mock_vector_store = MagicMock() | |
| # 2. Define fake documents to be returned | |
| mock_doc = MagicMock() | |
| mock_doc.page_content = "Pikachu is a mouse Pokémon." | |
| mock_vector_store.similarity_search.return_value = [mock_doc] | |
| # 3. Call the function | |
| query = "Who is Pikachu?" | |
| results = find_similar_documents(mock_vector_store, query, k=1) | |
| # 4. Assertions | |
| # Check that similarity_search was called with the right arguments | |
| mock_vector_store.similarity_search.assert_called_once_with(query, k=1) | |
| # Check the return value | |
| assert len(results) == 1 | |
| assert results[0].page_content == "Pikachu is a mouse Pokémon." | |
| def test_create_prompt(): | |
| # 1. Setup mock context (list of Document-like objects) | |
| mock_doc1 = MagicMock() | |
| mock_doc1.page_content = "Rayquaza lives in the Sky Pillar." | |
| mock_doc2 = MagicMock() | |
| mock_doc2.page_content = "Use a Master Ball for a 100% catch rate." | |
| context = [mock_doc1, mock_doc2] | |
| message = "Where can I find Rayquaza?" | |
| history = [] # Currently unused in your helpers.py but required by signature | |
| # 2. Call the function | |
| prompt = create_prompt(message, history, context) | |
| # 3. Assertions | |
| # Check if the persona is present | |
| assert "Pokémon Emerald Research Assistant" in prompt | |
| # Check if both context pieces are joined by the separator "---" | |
| assert "Rayquaza lives in the Sky Pillar." in prompt | |
| assert "---" in prompt | |
| assert "Use a Master Ball for a 100% catch rate." in prompt | |
| # Check if the user message is included | |
| assert "Where can I find Rayquaza?" in prompt | |
| # Check if Pydantic format instructions are present | |
| # (The parser adds specific JSON schema text) | |
| assert "json" in prompt.lower() | |
| assert "properties" in prompt.lower() |