From a2e0f2cefd50fabb9f3beac535913feea927ce3f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 12 Jul 2025 23:00:49 +0300 Subject: [PATCH] Remove Claude nonsense, use context manager --- src/haiku/rag/client.py | 27 +- tests/test_client.py | 638 +++++++++++++++++++--------------------- tests/test_monitor.py | 20 +- tests/test_rebuild.py | 71 +++-- tests/test_settings.py | 42 ++- 5 files changed, 366 insertions(+), 432 deletions(-) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 72e5b0a6..ca30098c 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -166,29 +166,26 @@ class HaikuRAG: # Create a temporary file with the appropriate extension with tempfile.NamedTemporaryFile( - mode="wb", suffix=file_extension, delete=False + mode="wb", suffix=file_extension ) as temp_file: temp_file.write(response.content) + temp_file.flush() # Ensure content is written to disk temp_path = Path(temp_file.name) - try: # Parse the content using FileReader content = FileReader.parse_file(temp_path) - # Merge metadata with contentType and md5 - metadata.update({"contentType": content_type, "md5": md5_hash}) + # Merge metadata with contentType and md5 + metadata.update({"contentType": content_type, "md5": md5_hash}) - if existing_doc: - existing_doc.content = content - existing_doc.metadata = metadata - return await self.update_document(existing_doc) - else: - return await self.create_document( - content=content, uri=url, metadata=metadata - ) - finally: - # Clean up temporary file - temp_path.unlink(missing_ok=True) + if existing_doc: + existing_doc.content = content + existing_doc.metadata = metadata + return await self.update_document(existing_doc) + else: + return await self.create_document( + content=content, uri=url, metadata=metadata + ) def _get_extension_from_content_type_or_url( self, url: str, content_type: str diff --git a/tests/test_client.py b/tests/test_client.py index 086facc0..c5aaeb1c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -12,300 +12,270 @@ from haiku.rag.client import HaikuRAG @pytest.mark.asyncio async def test_client_document_crud(qa_corpus: Dataset): """Test HaikuRAG CRUD operations for documents.""" - # Create client with in-memory database - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Get test data + first_doc = qa_corpus[0] + document_text = first_doc["document_extracted"] + test_uri = "file:///path/to/test.txt" + test_metadata = {"source": "test", "topic": "testing"} - # Get test data - first_doc = qa_corpus[0] - document_text = first_doc["document_extracted"] - test_uri = "file:///path/to/test.txt" - test_metadata = {"source": "test", "topic": "testing"} + # Test create_document + created_doc = await client.create_document( + content=document_text, uri=test_uri, metadata=test_metadata + ) - # Test create_document - created_doc = await client.create_document( - content=document_text, uri=test_uri, metadata=test_metadata - ) + assert created_doc.id is not None + assert created_doc.content == document_text + assert created_doc.uri == test_uri + assert created_doc.metadata == test_metadata - assert created_doc.id is not None - assert created_doc.content == document_text - assert created_doc.uri == test_uri - assert created_doc.metadata == test_metadata + # Test get_document_by_id + retrieved_doc = await client.get_document_by_id(created_doc.id) + assert retrieved_doc is not None + assert retrieved_doc.id == created_doc.id + assert retrieved_doc.content == document_text + assert retrieved_doc.uri == test_uri - # Test get_document_by_id - retrieved_doc = await client.get_document_by_id(created_doc.id) - assert retrieved_doc is not None - assert retrieved_doc.id == created_doc.id - assert retrieved_doc.content == document_text - assert retrieved_doc.uri == test_uri + # Test get_document_by_uri + retrieved_by_uri = await client.get_document_by_uri(test_uri) + assert retrieved_by_uri is not None + assert retrieved_by_uri.id == created_doc.id + assert retrieved_by_uri.content == document_text - # Test get_document_by_uri - retrieved_by_uri = await client.get_document_by_uri(test_uri) - assert retrieved_by_uri is not None - assert retrieved_by_uri.id == created_doc.id - assert retrieved_by_uri.content == document_text + # Test get_document_by_uri with non-existent URI + non_existent = await client.get_document_by_uri("file:///non/existent.txt") + assert non_existent is None - # Test get_document_by_uri with non-existent URI - non_existent = await client.get_document_by_uri("file:///non/existent.txt") - assert non_existent is None + # Test update_document + retrieved_doc.content = "Updated content" + retrieved_doc.uri = "file:///updated/path.txt" + updated_doc = await client.update_document(retrieved_doc) + assert updated_doc.content == "Updated content" + assert updated_doc.uri == "file:///updated/path.txt" - # Test update_document - retrieved_doc.content = "Updated content" - retrieved_doc.uri = "file:///updated/path.txt" - updated_doc = await client.update_document(retrieved_doc) - assert updated_doc.content == "Updated content" - assert updated_doc.uri == "file:///updated/path.txt" + # Test list_documents + all_docs = await client.list_documents() + assert len(all_docs) == 1 + assert all_docs[0].id == created_doc.id - # Test list_documents - all_docs = await client.list_documents() - assert len(all_docs) == 1 - assert all_docs[0].id == created_doc.id + # Test list_documents with pagination + limited_docs = await client.list_documents(limit=10, offset=0) + assert len(limited_docs) == 1 - # Test list_documents with pagination - limited_docs = await client.list_documents(limit=10, offset=0) - assert len(limited_docs) == 1 + # Test delete_document + deleted = await client.delete_document(created_doc.id) + assert deleted is True - # Test delete_document - deleted = await client.delete_document(created_doc.id) - assert deleted is True + # Verify document is gone + retrieved_doc = await client.get_document_by_id(created_doc.id) + assert retrieved_doc is None - # Verify document is gone - retrieved_doc = await client.get_document_by_id(created_doc.id) - assert retrieved_doc is None - - # Test delete non-existent document - deleted_again = await client.delete_document(created_doc.id) - assert deleted_again is False - - client.close() + # Test delete non-existent document + deleted_again = await client.delete_document(created_doc.id) + assert deleted_again is False @pytest.mark.asyncio async def test_client_create_document_from_source(): """Test creating a document from a file source.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + with tempfile.TemporaryDirectory() as temp_dir: + test_content = "This is test content from a file." + temp_path = Path(temp_dir) / "test.txt" + temp_path.write_text(test_content) - # Create a temporary text file - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - test_content = "This is test content from a file." - f.write(test_content) - temp_path = Path(f.name) + # Test create_document_from_source with Path + doc = await client.create_document_from_source(source=temp_path) - try: - # Test create_document_from_source with Path - doc = await client.create_document_from_source( - source=temp_path, metadata={"source_type": "file"} - ) + assert doc.id is not None + assert doc.content == test_content + assert doc.uri == temp_path.as_uri() + assert "contentType" in doc.metadata + assert "md5" in doc.metadata + assert doc.metadata["contentType"] == "text/plain" - assert doc.id is not None - assert doc.content == test_content - assert doc.uri == temp_path.as_uri() - assert doc.metadata["source_type"] == "file" - assert "contentType" in doc.metadata - assert "md5" in doc.metadata - assert doc.metadata["contentType"] == "text/plain" + # Test create_document_from_source with string path + doc2 = await client.create_document_from_source(source=str(temp_path)) - # Test create_document_from_source with string path - doc2 = await client.create_document_from_source(source=str(temp_path)) - - assert doc2.id is not None - assert doc2.content == test_content - assert doc2.uri == temp_path.as_uri() - assert "contentType" in doc2.metadata - assert "md5" in doc2.metadata - - finally: - # Clean up - temp_path.unlink() - client.close() + assert doc2.id is not None + assert doc2.content == test_content + assert doc2.uri == temp_path.as_uri() + assert "contentType" in doc2.metadata + assert "md5" in doc2.metadata @pytest.mark.asyncio async def test_client_create_document_from_source_unsupported(): """Test creating a document from an unsupported file type.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Create a temporary file with unsupported extension + with tempfile.NamedTemporaryFile( + mode="w", suffix=".unsupported", delete=False + ) as f: + f.write("content") + temp_path = Path(f.name) - # Create a temporary file with unsupported extension - with tempfile.NamedTemporaryFile( - mode="w", suffix=".unsupported", delete=False - ) as f: - f.write("content") - temp_path = Path(f.name) - - try: - # Should raise ValueError for unsupported extension - with pytest.raises(ValueError, match="Unsupported file extension"): - await client.create_document_from_source(temp_path) - - finally: - temp_path.unlink() - client.close() + # Should raise ValueError for unsupported extension + with pytest.raises(ValueError, match="Unsupported file extension"): + await client.create_document_from_source(temp_path) @pytest.mark.asyncio async def test_client_create_document_from_source_nonexistent(): """Test creating a document from a non-existent file.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + non_existent_path = Path("/non/existent/file.txt") - non_existent_path = Path("/non/existent/file.txt") - - # Should raise ValueError when file doesn't exist - with pytest.raises(ValueError, match="File does not exist"): - await client.create_document_from_source(non_existent_path) - - client.close() + # Should raise ValueError when file doesn't exist + with pytest.raises(ValueError, match="File does not exist"): + await client.create_document_from_source(non_existent_path) @pytest.mark.asyncio async def test_client_create_document_from_url(): """Test creating a document from a URL.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Mock the HTTP response + mock_response = AsyncMock() + mock_response.content = b"

Test Page

This is test content from a webpage.

" + mock_response.headers = {"content-type": "text/html"} + mock_response.raise_for_status = AsyncMock() - # Mock the HTTP response - mock_response = AsyncMock() - mock_response.content = b"

Test Page

This is test content from a webpage.

" - mock_response.headers = {"content-type": "text/html"} - mock_response.raise_for_status = AsyncMock() + with patch("httpx.AsyncClient.get", return_value=mock_response): + doc = await client.create_document_from_source( + source="https://example.com/test.html", metadata={"source_type": "web"} + ) - with patch("httpx.AsyncClient.get", return_value=mock_response): - doc = await client.create_document_from_source( - source="https://example.com/test.html", metadata={"source_type": "web"} - ) - - assert doc.id is not None - assert "Test Page" in doc.content - assert "test content" in doc.content - assert doc.uri == "https://example.com/test.html" - assert doc.metadata["source_type"] == "web" - assert "contentType" in doc.metadata - assert "md5" in doc.metadata - assert doc.metadata["contentType"] == "text/html" - - client.close() + assert doc.id is not None + assert "Test Page" in doc.content + assert "test content" in doc.content + assert doc.uri == "https://example.com/test.html" + assert doc.metadata["source_type"] == "web" + assert "contentType" in doc.metadata + assert "md5" in doc.metadata + assert doc.metadata["contentType"] == "text/html" @pytest.mark.asyncio async def test_client_create_document_from_url_with_different_content_types(): """Test creating documents from URLs with different content types.""" - client = HaikuRAG(":memory:") - - # Test JSON content - mock_json_response = AsyncMock() - mock_json_response.content = ( - b'{"title": "Test JSON", "content": "This is JSON content"}' - ) - mock_json_response.headers = {"content-type": "application/json"} - mock_json_response.raise_for_status = AsyncMock() - - with patch("httpx.AsyncClient.get", return_value=mock_json_response): - doc = await client.create_document_from_source( - "https://api.example.com/data.json" + async with HaikuRAG(":memory:") as client: + # Test JSON content + mock_json_response = AsyncMock() + mock_json_response.content = ( + b'{"title": "Test JSON", "content": "This is JSON content"}' ) + mock_json_response.headers = {"content-type": "application/json"} + mock_json_response.raise_for_status = AsyncMock() - assert doc.id is not None - assert "Test JSON" in doc.content - assert doc.uri == "https://api.example.com/data.json" - assert "contentType" in doc.metadata - assert "md5" in doc.metadata - assert doc.metadata["contentType"] == "application/json" + with patch("httpx.AsyncClient.get", return_value=mock_json_response): + doc = await client.create_document_from_source( + "https://api.example.com/data.json" + ) - # Test plain text content - mock_text_response = AsyncMock() - mock_text_response.content = b"This is plain text content from a URL." - mock_text_response.headers = {"content-type": "text/plain"} - mock_text_response.raise_for_status = AsyncMock() + assert doc.id is not None + assert "Test JSON" in doc.content + assert doc.uri == "https://api.example.com/data.json" + assert "contentType" in doc.metadata + assert "md5" in doc.metadata + assert doc.metadata["contentType"] == "application/json" - with patch("httpx.AsyncClient.get", return_value=mock_text_response): - doc = await client.create_document_from_source("https://example.com/readme.txt") + # Test plain text content + mock_text_response = AsyncMock() + mock_text_response.content = b"This is plain text content from a URL." + mock_text_response.headers = {"content-type": "text/plain"} + mock_text_response.raise_for_status = AsyncMock() - assert doc.id is not None - assert doc.content == "This is plain text content from a URL." - assert doc.uri == "https://example.com/readme.txt" - assert "contentType" in doc.metadata - assert "md5" in doc.metadata - assert doc.metadata["contentType"] == "text/plain" + with patch("httpx.AsyncClient.get", return_value=mock_text_response): + doc = await client.create_document_from_source( + "https://example.com/readme.txt" + ) - client.close() + assert doc.id is not None + assert doc.content == "This is plain text content from a URL." + assert doc.uri == "https://example.com/readme.txt" + assert "contentType" in doc.metadata + assert "md5" in doc.metadata + assert doc.metadata["contentType"] == "text/plain" @pytest.mark.asyncio async def test_client_create_document_from_url_unsupported_content(): """Test creating a document from URL with unsupported content type.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Mock response with unsupported content type + mock_response = AsyncMock() + mock_response.content = b"binary content" + mock_response.headers = {"content-type": "application/octet-stream"} + mock_response.raise_for_status = AsyncMock() - # Mock response with unsupported content type - mock_response = AsyncMock() - mock_response.content = b"binary content" - mock_response.headers = {"content-type": "application/octet-stream"} - mock_response.raise_for_status = AsyncMock() - - with patch("httpx.AsyncClient.get", return_value=mock_response): - with pytest.raises(ValueError, match="Unsupported content type"): - await client.create_document_from_source("https://example.com/binary.bin") - - client.close() + with patch("httpx.AsyncClient.get", return_value=mock_response): + with pytest.raises(ValueError, match="Unsupported content type"): + await client.create_document_from_source( + "https://example.com/binary.bin" + ) @pytest.mark.asyncio async def test_client_create_document_from_url_http_error(): """Test handling HTTP errors when creating document from URL.""" - client = HaikuRAG(":memory:") - - with patch("httpx.AsyncClient.get") as mock_get: - mock_get.side_effect = httpx.HTTPStatusError( - "404 Not Found", - request=httpx.Request("GET", "https://example.com/notfound.html"), - response=httpx.Response(404), - ) - - with pytest.raises(httpx.HTTPStatusError): - await client.create_document_from_source( - "https://example.com/notfound.html" + async with HaikuRAG(":memory:") as client: + with patch("httpx.AsyncClient.get") as mock_get: + mock_get.side_effect = httpx.HTTPStatusError( + "404 Not Found", + request=httpx.Request("GET", "https://example.com/notfound.html"), + response=httpx.Response(404), ) - client.close() + with pytest.raises(httpx.HTTPStatusError): + await client.create_document_from_source( + "https://example.com/notfound.html" + ) @pytest.mark.asyncio async def test_get_extension_from_content_type_or_url(): """Test the helper method for determining file extensions.""" - client = HaikuRAG(":memory:") - - # Test content type mappings - assert client._get_extension_from_content_type_or_url("", "text/html") == ".html" - assert ( - client._get_extension_from_content_type_or_url("", "application/pdf") == ".pdf" - ) - assert client._get_extension_from_content_type_or_url("", "text/plain") == ".txt" - - # Test URL extension detection - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/doc.pdf", "" + async with HaikuRAG(":memory:") as client: + # Test content type mappings + assert ( + client._get_extension_from_content_type_or_url("", "text/html") == ".html" ) - == ".pdf" - ) - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/data.json", "" + assert ( + client._get_extension_from_content_type_or_url("", "application/pdf") + == ".pdf" ) - == ".json" - ) - - # Test default fallback - assert ( - client._get_extension_from_content_type_or_url("https://example.com/", "") - == ".html" - ) - - # Test content type priority over URL extension - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/file.txt", "application/pdf" + assert ( + client._get_extension_from_content_type_or_url("", "text/plain") == ".txt" ) - == ".pdf" - ) - client.close() + # Test URL extension detection + assert ( + client._get_extension_from_content_type_or_url( + "https://example.com/doc.pdf", "" + ) + == ".pdf" + ) + assert ( + client._get_extension_from_content_type_or_url( + "https://example.com/data.json", "" + ) + == ".json" + ) + + # Test default fallback + assert ( + client._get_extension_from_content_type_or_url("https://example.com/", "") + == ".html" + ) + + # Test content type priority over URL extension + assert ( + client._get_extension_from_content_type_or_url( + "https://example.com/file.txt", "application/pdf" + ) + == ".pdf" + ) @pytest.mark.asyncio @@ -313,165 +283,147 @@ async def test_client_metadata_content_type_and_md5(): """Test that contentType and md5 metadata are correctly set.""" import hashlib - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Create a temporary file with known content + test_content = "Test content for MD5 calculation." + expected_md5 = hashlib.md5(test_content.encode()).hexdigest() - # Create a temporary file with known content - test_content = "Test content for MD5 calculation." - expected_md5 = hashlib.md5(test_content.encode()).hexdigest() + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / "test.txt" + temp_path.write_text(test_content) - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - f.write(test_content) - temp_path = Path(f.name) + doc = await client.create_document_from_source(temp_path) - try: - doc = await client.create_document_from_source(temp_path) + assert doc.metadata["contentType"] == "text/plain" + assert doc.metadata["md5"] == expected_md5 - assert doc.metadata["contentType"] == "text/plain" - assert doc.metadata["md5"] == expected_md5 + mock_response = AsyncMock() + mock_response.content = test_content.encode() + mock_response.headers = {"content-type": "text/plain"} + mock_response.raise_for_status = AsyncMock() - mock_response = AsyncMock() - mock_response.content = test_content.encode() - mock_response.headers = {"content-type": "text/plain"} - mock_response.raise_for_status = AsyncMock() + with patch("httpx.AsyncClient.get", return_value=mock_response): + url_doc = await client.create_document_from_source( + "https://example.com/test.txt" + ) - with patch("httpx.AsyncClient.get", return_value=mock_response): - url_doc = await client.create_document_from_source( - "https://example.com/test.txt" - ) - - assert url_doc.metadata["contentType"] == "text/plain" - assert url_doc.metadata["md5"] == expected_md5 - - finally: - temp_path.unlink() - client.close() + assert url_doc.metadata["contentType"] == "text/plain" + assert url_doc.metadata["md5"] == expected_md5 @pytest.mark.asyncio async def test_client_create_update_no_op_behavior(): """Test create/update/no-op behavior based on MD5 changes.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Create a temporary file + test_content = "Original content for testing." + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / "test.txt" + temp_path.write_text(test_content) - # Create a temporary file - test_content = "Original content for testing." - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - f.write(test_content) - temp_path = Path(f.name) + # First call - should create new document + doc1 = await client.create_document_from_source(temp_path) + assert doc1.id is not None + assert doc1.content == test_content + original_id = doc1.id - try: - # First call - should create new document - doc1 = await client.create_document_from_source(temp_path) - assert doc1.id is not None - assert doc1.content == test_content - original_id = doc1.id + # Second call with same content - should return existing document (no-op) + doc2 = await client.create_document_from_source(temp_path) + assert doc2.id == original_id # Same document + assert doc2.content == test_content - # Second call with same content - should return existing document (no-op) - doc2 = await client.create_document_from_source(temp_path) - assert doc2.id == original_id # Same document - assert doc2.content == test_content + # Modify file content + updated_content = "Updated content for testing." + temp_path.write_text(updated_content) - # Modify file content - updated_content = "Updated content for testing." - temp_path.write_text(updated_content) + # Third call with changed content - should update existing document + doc3 = await client.create_document_from_source(temp_path) + assert doc3.id == original_id # Same document ID + assert doc3.content == updated_content # Updated content - # Third call with changed content - should update existing document - doc3 = await client.create_document_from_source(temp_path) - assert doc3.id == original_id # Same document ID - assert doc3.content == updated_content # Updated content - - # Verify the document was actually updated in database - retrieved_doc = await client.get_document_by_id(original_id) - assert retrieved_doc is not None - assert retrieved_doc.content == updated_content - - finally: - temp_path.unlink() - client.close() + # Verify the document was actually updated in database + retrieved_doc = await client.get_document_by_id(original_id) + assert retrieved_doc is not None + assert retrieved_doc.content == updated_content @pytest.mark.asyncio async def test_client_url_create_update_no_op_behavior(): """Test create/update/no-op behavior for URLs based on MD5 changes.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + url = "https://example.com/test.txt" + original_content = b"Original URL content" + updated_content = b"Updated URL content" - url = "https://example.com/test.txt" - original_content = b"Original URL content" - updated_content = b"Updated URL content" + # Mock first response + mock_response1 = AsyncMock() + mock_response1.content = original_content + mock_response1.headers = {"content-type": "text/plain"} + mock_response1.raise_for_status = AsyncMock() - # Mock first response - mock_response1 = AsyncMock() - mock_response1.content = original_content - mock_response1.headers = {"content-type": "text/plain"} - mock_response1.raise_for_status = AsyncMock() + with patch("httpx.AsyncClient.get", return_value=mock_response1): + # First call - should create new document + doc1 = await client.create_document_from_source(url) + assert doc1.id is not None + original_id = doc1.id - with patch("httpx.AsyncClient.get", return_value=mock_response1): - # First call - should create new document - doc1 = await client.create_document_from_source(url) - assert doc1.id is not None - original_id = doc1.id + # Second call with same content - should return existing document (no-op) + doc2 = await client.create_document_from_source(url) + assert doc2.id == original_id # Same document - # Second call with same content - should return existing document (no-op) - doc2 = await client.create_document_from_source(url) - assert doc2.id == original_id # Same document + mock_response2 = AsyncMock() + mock_response2.content = updated_content + mock_response2.headers = {"content-type": "text/plain"} + mock_response2.raise_for_status = AsyncMock() - mock_response2 = AsyncMock() - mock_response2.content = updated_content - mock_response2.headers = {"content-type": "text/plain"} - mock_response2.raise_for_status = AsyncMock() - - with patch("httpx.AsyncClient.get", return_value=mock_response2): - # Third call with changed content - should update existing document - doc3 = await client.create_document_from_source(url) - assert doc3.id == original_id # Same document ID - assert doc3.content == updated_content.decode() # Updated content - - client.close() + with patch("httpx.AsyncClient.get", return_value=mock_response2): + # Third call with changed content - should update existing document + doc3 = await client.create_document_from_source(url) + assert doc3.id == original_id # Same document ID + assert doc3.content == updated_content.decode() # Updated content @pytest.mark.asyncio async def test_client_search(): """Test HaikuRAG search functionality.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + # Add multiple documents to search from + doc1_text = "Python is a high-level programming language known for its simplicity and readability." + doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming." + doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights." - # Add multiple documents to search from - doc1_text = "Python is a high-level programming language known for its simplicity and readability." - doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming." - doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights." + # Create documents + doc1 = await client.create_document( + content=doc1_text, uri="doc1.txt", metadata={"topic": "python"} + ) + doc2 = await client.create_document( + content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"} + ) + await client.create_document( + content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"} + ) - # Create documents - doc1 = await client.create_document( - content=doc1_text, uri="doc1.txt", metadata={"topic": "python"} - ) - doc2 = await client.create_document( - content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"} - ) - await client.create_document( - content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"} - ) + # Test search with keyword that should match doc1 + results = await client.search("Python programming", limit=3) - # Test search with keyword that should match doc1 - results = await client.search("Python programming", limit=3) + assert len(results) > 0 + assert all(len(result) == 2 for result in results) - assert len(results) > 0 - assert all(len(result) == 2 for result in results) + # Verify first result is from the Python document (doc1) + first_chunk, _ = results[0] + assert first_chunk.document_id == doc1.id - # Verify first result is from the Python document (doc1) - first_chunk, _ = results[0] - assert first_chunk.document_id == doc1.id + # Test search with different query + ml_results = await client.search("machine learning data", limit=2) + assert len(ml_results) > 0 - # Test search with different query - ml_results = await client.search("machine learning data", limit=2) - assert len(ml_results) > 0 + # Verify first result is from the machine learning document (doc2) + first_ml_chunk, _ = ml_results[0] + assert first_ml_chunk.document_id == doc2.id - # Verify first result is from the machine learning document (doc2) - first_ml_chunk, _ = ml_results[0] - assert first_ml_chunk.document_id == doc2.id - - # Test search with limit parameter - limited_results = await client.search("programming", limit=1) - assert len(limited_results) <= 1 - - client.close() + # Test search with limit parameter + limited_results = await client.search("programming", limit=1) + assert len(limited_results) <= 1 @pytest.mark.asyncio diff --git a/tests/test_monitor.py b/tests/test_monitor.py index ac909631..73992416 100644 --- a/tests/test_monitor.py +++ b/tests/test_monitor.py @@ -13,11 +13,10 @@ from haiku.rag.store.models.document import Document async def test_file_watcher_upsert_document(): """Test FileWatcher._upsert_document method.""" - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - f.write("Test content for file watcher") - temp_path = Path(f.name) + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / "test.txt" + temp_path.write_text("Test content for file watcher") - try: mock_client = AsyncMock(spec=HaikuRAG) mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri()) mock_client.create_document_from_source.return_value = mock_doc @@ -32,19 +31,15 @@ async def test_file_watcher_upsert_document(): mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) mock_client.create_document_from_source.assert_called_once_with(str(temp_path)) - finally: - temp_path.unlink(missing_ok=True) - @pytest.mark.asyncio async def test_file_watcher_upsert_existing_document(): """Test FileWatcher._upsert_document with existing document.""" - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - f.write("Test content for file watcher") - temp_path = Path(f.name) + with tempfile.TemporaryDirectory() as temp_dir: + temp_path = Path(temp_dir) / "test.txt" + temp_path.write_text("Test content for file watcher") - try: mock_client = AsyncMock(spec=HaikuRAG) existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri()) updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri()) @@ -61,9 +56,6 @@ async def test_file_watcher_upsert_existing_document(): mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri()) mock_client.create_document_from_source.assert_called_once_with(str(temp_path)) - finally: - temp_path.unlink(missing_ok=True) - @pytest.mark.asyncio async def test_file_watcher_delete_document(): diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 3254ce1d..04fd8533 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -8,45 +8,42 @@ from haiku.rag.store.models.document import Document @pytest.mark.asyncio async def test_rebuild_database(qa_corpus: Dataset): """Test rebuild functionality with existing documents.""" - client = HaikuRAG(":memory:") + async with HaikuRAG(":memory:") as client: + created_docs: list[Document] = [] + for content in qa_corpus["document_extracted"][:3]: + doc = await client.create_document( + content=content, + ) + created_docs.append(doc) - created_docs: list[Document] = [] - for content in qa_corpus["document_extracted"][:3]: - doc = await client.create_document( - content=content, - ) - created_docs.append(doc) + documents_before = await client.list_documents() + assert len(documents_before) == 3 - documents_before = await client.list_documents() - assert len(documents_before) == 3 - - chunks_before = [] - for doc in created_docs: - assert doc.id is not None - doc_chunks = await client.chunk_repository.get_by_document_id(doc.id) - chunks_before.extend(doc_chunks) - - assert len(chunks_before) > 0 - - # Perform rebuild - processed_doc_ids = [] - async for doc_id in client.rebuild_database(): - processed_doc_ids.append(doc_id) - - # Verify all documents were processed - expected_doc_ids = [doc.id for doc in created_docs] - assert set(processed_doc_ids) == set(expected_doc_ids) - - documents_after = await client.list_documents() - assert len(documents_after) == 3 - - # Verify chunks were recreated - chunks_after = [] - for doc in documents_after: - if doc.id is not None: + chunks_before = [] + for doc in created_docs: + assert doc.id is not None doc_chunks = await client.chunk_repository.get_by_document_id(doc.id) - chunks_after.extend(doc_chunks) + chunks_before.extend(doc_chunks) - assert len(chunks_after) > 0 + assert len(chunks_before) > 0 - client.close() + # Perform rebuild + processed_doc_ids = [] + async for doc_id in client.rebuild_database(): + processed_doc_ids.append(doc_id) + + # Verify all documents were processed + expected_doc_ids = [doc.id for doc in created_docs] + assert set(processed_doc_ids) == set(expected_doc_ids) + + documents_after = await client.list_documents() + assert len(documents_after) == 3 + + # Verify chunks were recreated + chunks_after = [] + for doc in documents_after: + if doc.id is not None: + doc_chunks = await client.chunk_repository.get_by_document_id(doc.id) + chunks_after.extend(doc_chunks) + + assert len(chunks_after) > 0 diff --git a/tests/test_settings.py b/tests/test_settings.py index 0db16cc8..6a531685 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -45,10 +45,9 @@ def test_settings_save_and_retrieve(): async def test_config_validation_on_db_load(): """Test that config validation fails when loading db with mismatched settings.""" # Create a temporary database file - with tempfile.NamedTemporaryFile(suffix=".sqlite", delete=False) as tmp: + with tempfile.NamedTemporaryFile(suffix=".sqlite") as tmp: db_path = Path(tmp.name) - try: # Create store and save settings store1 = Store(db_path) SettingsRepository(store1) @@ -58,28 +57,25 @@ async def test_config_validation_on_db_load(): original_chunk_size = Config.CHUNK_SIZE Config.CHUNK_SIZE = 999 - # Loading the database should raise ConfigMismatchError - with pytest.raises(ConfigMismatchError) as exc_info: - Store(db_path) + try: + # Loading the database should raise ConfigMismatchError + with pytest.raises(ConfigMismatchError) as exc_info: + Store(db_path) - assert "CHUNK_SIZE" in str(exc_info.value) - assert "Consider rebuilding" in str(exc_info.value) + assert "CHUNK_SIZE" in str(exc_info.value) + assert "Consider rebuilding" in str(exc_info.value) - # Rebuild - async with HaikuRAG(db_path=db_path, skip_validation=True) as client: - async for _ in client.rebuild_database(): - pass # Process all documents + # Rebuild + async with HaikuRAG(db_path=db_path, skip_validation=True) as client: + async for _ in client.rebuild_database(): + pass # Process all documents - # Verify we can now load the database without exception (settings were updated) - store2 = Store(db_path) - settings_repo2 = SettingsRepository(store2) - db_settings = settings_repo2.get() - assert db_settings["CHUNK_SIZE"] == 999 - store2.close() + # Verify we can now load the database without exception (settings were updated) + store2 = Store(db_path) + settings_repo2 = SettingsRepository(store2) + db_settings = settings_repo2.get() + assert db_settings["CHUNK_SIZE"] == 999 + store2.close() - Config.CHUNK_SIZE = original_chunk_size - - finally: - # Cleanup - if db_path.exists(): - db_path.unlink() + finally: + Config.CHUNK_SIZE = original_chunk_size