feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

tests/test_agent.py

@@ -0,0 +1,128 @@
+import os
+import sys
+import json
+import tempfile
+import shutil
+from pathlib import Path
+from unittest import mock
+
+# Ensure the environment variable is set before importing the module
+os.environ["OPENAI_API_KEY"] = "test-key"
+
+# Import the module after setting the environment variable
+import src.index as index
+
+# Helper to create a temporary text file
+def create_temp_txt(folder: Path, name: str, content: str):
+    file_path = folder / name
+    file_path.write_text(content, encoding="utf-8")
+    return file_path
+
+def test_split_text_basic():
+    text = "Sentence one. Sentence two. Sentence three."
+    chunks = index._split_text(text, max_chunk_size=50)
+    assert len(chunks) == 3
+    assert chunks[0] == "Sentence one."
+    assert chunks[1] == "Sentence two."
+    assert chunks[2] == "Sentence three."
+
+def test_split_text_long_sentence():
+    long_sentence = "A" * 200
+    text = f"{long_sentence}. Another short sentence."
+    chunks = index._split_text(text, max_chunk_size=100)
+    # The long sentence should be split into two chunks
+    assert len(chunks) == 2
+    assert chunks[0].startswith("A" * 100)
+    assert chunks[1].startswith("A" * 100)
+
+def test_ingest_folder(monkeypatch):
+    # Create a temporary directory with a single .txt file
+    temp_dir = Path(tempfile.mkdtemp())
+    try:
+        content = "Hello world. This is a test."
+        create_temp_txt(temp_dir, "test.txt", content)
+
+        # Mock the add_documents method to capture its arguments
+        captured = {}
+        def mock_add_documents(self, documents, ids):
+            captured["documents"] = documents
+            captured["ids"] = ids
+
+        monkeypatch.setattr(index.ChromaDBWrapper, "add_documents", mock_add_documents)
+
+        # Instantiate the agent with a dummy db wrapper
+        dummy_db = index.ChromaDBWrapper()
+        agent = index.RAGAgent(dummy_db)
+
+        # Run ingestion
+        agent.ingest_folder(str(temp_dir))
+
+        # Verify that documents were split and added
+        assert "documents" in captured
+        assert "ids" in captured
+        assert len(captured["documents"]) == 2  # two sentences
+        assert captured["documents"][0] == "Hello world."
+        assert captured["documents"][1] == "This is a test."
+        assert len(captured["ids"]) == 2
+        assert captured["ids"][0].startswith("test_")
+    finally:
+        shutil.rmtree(temp_dir)
+
+def test_answer_query_local(monkeypatch):
+    # Dummy database that returns a relevant document
+    class DummyDB:
+        def query(self, query_text, k=5):
+            return [("Relevant context about Python.", 0.8)]
+
+    dummy_db = DummyDB()
+    agent = index.RAGAgent(dummy_db)
+
+    # Mock the OpenAI ChatCompletion to return a predictable answer
+    mock_response = {
+        "choices": [
+            {"message": {"content": "Python is a programming language."}}
+        ]
+    }
+    monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)
+
+    answer = agent.answer_query("What is Python?")
+    assert answer == "Python is a programming language."
+
+def test_answer_query_fallback(monkeypatch):
+    # Dummy database that returns no relevant documents
+    class DummyDB:
+        def query(self, query_text, k=5):
+            return []
+
+    dummy_db = DummyDB()
+    agent = index.RAGAgent(dummy_db)
+
+    # Mock web_search to return snippets
+    monkeypatch.setattr(index, "web_search", lambda query, max_results=3: ["Snippet about AI.", "Another snippet."])
+
+    # Mock the OpenAI ChatCompletion to return a predictable answer
+    mock_response = {
+        "choices": [
+            {"message": {"content": "AI stands for Artificial Intelligence."}}
+        ]
+    }
+    monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)
+
+    answer = agent.answer_query("What does AI stand for?")
+    assert answer == "AI stands for Artificial Intelligence."
+
+def test_web_search_mock(monkeypatch):
+    # Mock ddg to return predefined results
+    mock_results = [
+        {"body": "First snippet content."},
+        {"body": "Second snippet content."},
+    ]
+    monkeypatch.setattr(index.ddg, "__call__", lambda query, max_results=3: mock_results)
+
+    snippets = index.web_search("test query")
+    assert snippets == ["First snippet content.", "Second snippet content."]
+
+if __name__ == "__main__":
+    # Run tests manually if executed as a script
+    import pytest
+    sys.exit(pytest.main([__file__]))