Remove Claude nonsense, use context manager

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
+            # Merge metadata with contentType and md5
-                metadata.update({"contentType": content_type, "md5": md5_hash})
+            metadata.update({"contentType": content_type, "md5": md5_hash})
 
-                if existing_doc:
+            if existing_doc:
-                    existing_doc.content = content
+                existing_doc.content = content
-                    existing_doc.metadata = metadata
+                existing_doc.metadata = metadata
-                    return await self.update_document(existing_doc)
+                return await self.update_document(existing_doc)
-                else:
+            else:
-                    return await self.create_document(
+                return await self.create_document(
-                        content=content, uri=url, metadata=metadata
+                    content=content, uri=url, metadata=metadata
-                    )
+                )
-            finally:
-                # Clean up temporary file
-                temp_path.unlink(missing_ok=True)
 
     def _get_extension_from_content_type_or_url(
         self, url: str, content_type: str

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
+    async with HaikuRAG(":memory:") as client:
-    client = HaikuRAG(":memory:")
+        # 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
+        # Test create_document
-    first_doc = qa_corpus[0]
+        created_doc = await client.create_document(
-    document_text = first_doc["document_extracted"]
+            content=document_text, uri=test_uri, metadata=test_metadata
-    test_uri = "file:///path/to/test.txt"
+        )
-    test_metadata = {"source": "test", "topic": "testing"}
 
-    # Test create_document
+        assert created_doc.id is not None
-    created_doc = await client.create_document(
+        assert created_doc.content == document_text
-        content=document_text, uri=test_uri, metadata=test_metadata
+        assert created_doc.uri == test_uri
-    )
+        assert created_doc.metadata == test_metadata
 
-    assert created_doc.id is not None
+        # Test get_document_by_id
-    assert created_doc.content == document_text
+        retrieved_doc = await client.get_document_by_id(created_doc.id)
-    assert created_doc.uri == test_uri
+        assert retrieved_doc is not None
-    assert created_doc.metadata == test_metadata
+        assert retrieved_doc.id == created_doc.id
+        assert retrieved_doc.content == document_text
+        assert retrieved_doc.uri == test_uri
 
-    # Test get_document_by_id
+        # Test get_document_by_uri
-    retrieved_doc = await client.get_document_by_id(created_doc.id)
+        retrieved_by_uri = await client.get_document_by_uri(test_uri)
-    assert retrieved_doc is not None
+        assert retrieved_by_uri is not None
-    assert retrieved_doc.id == created_doc.id
+        assert retrieved_by_uri.id == created_doc.id
-    assert retrieved_doc.content == document_text
+        assert retrieved_by_uri.content == document_text
-    assert retrieved_doc.uri == test_uri
 
-    # Test get_document_by_uri
+        # Test get_document_by_uri with non-existent URI
-    retrieved_by_uri = await client.get_document_by_uri(test_uri)
+        non_existent = await client.get_document_by_uri("file:///non/existent.txt")
-    assert retrieved_by_uri is not None
+        assert non_existent is 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
+        # Test update_document
-    non_existent = await client.get_document_by_uri("file:///non/existent.txt")
+        retrieved_doc.content = "Updated content"
-    assert non_existent is None
+        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
+        # Test list_documents
-    retrieved_doc.content = "Updated content"
+        all_docs = await client.list_documents()
-    retrieved_doc.uri = "file:///updated/path.txt"
+        assert len(all_docs) == 1
-    updated_doc = await client.update_document(retrieved_doc)
+        assert all_docs[0].id == created_doc.id
-    assert updated_doc.content == "Updated content"
-    assert updated_doc.uri == "file:///updated/path.txt"
 
-    # Test list_documents
+        # Test list_documents with pagination
-    all_docs = await client.list_documents()
+        limited_docs = await client.list_documents(limit=10, offset=0)
-    assert len(all_docs) == 1
+        assert len(limited_docs) == 1
-    assert all_docs[0].id == created_doc.id
 
-    # Test list_documents with pagination
+        # Test delete_document
-    limited_docs = await client.list_documents(limit=10, offset=0)
+        deleted = await client.delete_document(created_doc.id)
-    assert len(limited_docs) == 1
+        assert deleted is True
 
-    # Test delete_document
+        # Verify document is gone
-    deleted = await client.delete_document(created_doc.id)
+        retrieved_doc = await client.get_document_by_id(created_doc.id)
-    assert deleted is True
+        assert retrieved_doc is None
 
-    # Verify document is gone
+        # Test delete non-existent document
-    retrieved_doc = await client.get_document_by_id(created_doc.id)
+        deleted_again = await client.delete_document(created_doc.id)
-    assert retrieved_doc is None
+        assert deleted_again is False
-
-    # Test delete non-existent document
-    deleted_again = await client.delete_document(created_doc.id)
-    assert deleted_again is False
-
-    client.close()
 
 
 @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
+            # Test create_document_from_source with Path
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+            doc = await client.create_document_from_source(source=temp_path)
-        test_content = "This is test content from a file."
-        f.write(test_content)
-        temp_path = Path(f.name)
 
-    try:
+            assert doc.id is not None
-        # Test create_document_from_source with Path
+            assert doc.content == test_content
-        doc = await client.create_document_from_source(
+            assert doc.uri == temp_path.as_uri()
-            source=temp_path, metadata={"source_type": "file"}
+            assert "contentType" in doc.metadata
-        )
+            assert "md5" in doc.metadata
+            assert doc.metadata["contentType"] == "text/plain"
 
-        assert doc.id is not None
+            # Test create_document_from_source with string path
-        assert doc.content == test_content
+            doc2 = await client.create_document_from_source(source=str(temp_path))
-        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
+            assert doc2.id is not None
-        doc2 = await client.create_document_from_source(source=str(temp_path))
+            assert doc2.content == test_content
-
+            assert doc2.uri == temp_path.as_uri()
-        assert doc2.id is not None
+            assert "contentType" in doc2.metadata
-        assert doc2.content == test_content
+            assert "md5" in doc2.metadata
-        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()
 
 
 @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
+            # Should raise ValueError for unsupported extension
-    with tempfile.NamedTemporaryFile(
+            with pytest.raises(ValueError, match="Unsupported file extension"):
-        mode="w", suffix=".unsupported", delete=False
+                await client.create_document_from_source(temp_path)
-    ) 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()
 
 
 @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"):
-    # Should raise ValueError when file doesn't exist
+            await client.create_document_from_source(non_existent_path)
-    with pytest.raises(ValueError, match="File does not exist"):
-        await client.create_document_from_source(non_existent_path)
-
-    client.close()
 
 
 @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"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
+        mock_response.headers = {"content-type": "text/html"}
+        mock_response.raise_for_status = AsyncMock()
 
-    # Mock the HTTP response
+        with patch("httpx.AsyncClient.get", return_value=mock_response):
-    mock_response = AsyncMock()
+            doc = await client.create_document_from_source(
-    mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
+                source="https://example.com/test.html", metadata={"source_type": "web"}
-    mock_response.headers = {"content-type": "text/html"}
+            )
-    mock_response.raise_for_status = AsyncMock()
 
-    with patch("httpx.AsyncClient.get", return_value=mock_response):
+            assert doc.id is not None
-        doc = await client.create_document_from_source(
+            assert "Test Page" in doc.content
-            source="https://example.com/test.html", metadata={"source_type": "web"}
+            assert "test content" in doc.content
-        )
+            assert doc.uri == "https://example.com/test.html"
-
+            assert doc.metadata["source_type"] == "web"
-        assert doc.id is not None
+            assert "contentType" in doc.metadata
-        assert "Test Page" in doc.content
+            assert "md5" in doc.metadata
-        assert "test content" in doc.content
+            assert doc.metadata["contentType"] == "text/html"
-        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()
 
 
 @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:")
+    async with HaikuRAG(":memory:") as client:
-
+        # Test JSON content
-    # Test JSON content
+        mock_json_response = AsyncMock()
-    mock_json_response = AsyncMock()
+        mock_json_response.content = (
-    mock_json_response.content = (
+            b'{"title": "Test JSON", "content": "This is JSON 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"
         )
+        mock_json_response.headers = {"content-type": "application/json"}
+        mock_json_response.raise_for_status = AsyncMock()
 
-        assert doc.id is not None
+        with patch("httpx.AsyncClient.get", return_value=mock_json_response):
-        assert "Test JSON" in doc.content
+            doc = await client.create_document_from_source(
-        assert doc.uri == "https://api.example.com/data.json"
+                "https://api.example.com/data.json"
-        assert "contentType" in doc.metadata
+            )
-        assert "md5" in doc.metadata
-        assert doc.metadata["contentType"] == "application/json"
 
-    # Test plain text content
+            assert doc.id is not None
-    mock_text_response = AsyncMock()
+            assert "Test JSON" in doc.content
-    mock_text_response.content = b"This is plain text content from a URL."
+            assert doc.uri == "https://api.example.com/data.json"
-    mock_text_response.headers = {"content-type": "text/plain"}
+            assert "contentType" in doc.metadata
-    mock_text_response.raise_for_status = AsyncMock()
+            assert "md5" in doc.metadata
+            assert doc.metadata["contentType"] == "application/json"
 
-    with patch("httpx.AsyncClient.get", return_value=mock_text_response):
+        # Test plain text content
-        doc = await client.create_document_from_source("https://example.com/readme.txt")
+        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
+        with patch("httpx.AsyncClient.get", return_value=mock_text_response):
-        assert doc.content == "This is plain text content from a URL."
+            doc = await client.create_document_from_source(
-        assert doc.uri == "https://example.com/readme.txt"
+                "https://example.com/readme.txt"
-        assert "contentType" in doc.metadata
+            )
-        assert "md5" in doc.metadata
-        assert doc.metadata["contentType"] == "text/plain"
 
-    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
+        with patch("httpx.AsyncClient.get", return_value=mock_response):
-    mock_response = AsyncMock()
+            with pytest.raises(ValueError, match="Unsupported content type"):
-    mock_response.content = b"binary content"
+                await client.create_document_from_source(
-    mock_response.headers = {"content-type": "application/octet-stream"}
+                    "https://example.com/binary.bin"
-    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()
 
 
 @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:")
+    async with HaikuRAG(":memory:") as client:
-
+        with patch("httpx.AsyncClient.get") as mock_get:
-    with patch("httpx.AsyncClient.get") as mock_get:
+            mock_get.side_effect = httpx.HTTPStatusError(
-        mock_get.side_effect = httpx.HTTPStatusError(
+                "404 Not Found",
-            "404 Not Found",
+                request=httpx.Request("GET", "https://example.com/notfound.html"),
-            request=httpx.Request("GET", "https://example.com/notfound.html"),
+                response=httpx.Response(404),
-            response=httpx.Response(404),
-        )
-
-        with pytest.raises(httpx.HTTPStatusError):
-            await client.create_document_from_source(
-                "https://example.com/notfound.html"
             )
 
-    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:")
+    async with HaikuRAG(":memory:") as client:
-
+        # Test content type mappings
-    # Test content type mappings
+        assert (
-    assert client._get_extension_from_content_type_or_url("", "text/html") == ".html"
+            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", ""
         )
-        == ".pdf"
+        assert (
-    )
+            client._get_extension_from_content_type_or_url("", "application/pdf")
-    assert (
+            == ".pdf"
-        client._get_extension_from_content_type_or_url(
-            "https://example.com/data.json", ""
         )
-        == ".json"
+        assert (
-    )
+            client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
-
-    # 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"
-    )
 
-    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
+        with tempfile.TemporaryDirectory() as temp_dir:
-    test_content = "Test content for MD5 calculation."
+            temp_path = Path(temp_dir) / "test.txt"
-    expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
+            temp_path.write_text(test_content)
 
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+            doc = await client.create_document_from_source(temp_path)
-        f.write(test_content)
-        temp_path = Path(f.name)
 
-    try:
+            assert doc.metadata["contentType"] == "text/plain"
-        doc = await client.create_document_from_source(temp_path)
+            assert doc.metadata["md5"] == expected_md5
 
-        assert doc.metadata["contentType"] == "text/plain"
+            mock_response = AsyncMock()
-        assert doc.metadata["md5"] == expected_md5
+            mock_response.content = test_content.encode()
+            mock_response.headers = {"content-type": "text/plain"}
+            mock_response.raise_for_status = AsyncMock()
 
-        mock_response = AsyncMock()
+            with patch("httpx.AsyncClient.get", return_value=mock_response):
-        mock_response.content = test_content.encode()
+                url_doc = await client.create_document_from_source(
-        mock_response.headers = {"content-type": "text/plain"}
+                    "https://example.com/test.txt"
-        mock_response.raise_for_status = AsyncMock()
+                )
 
-        with patch("httpx.AsyncClient.get", return_value=mock_response):
+                assert url_doc.metadata["contentType"] == "text/plain"
-            url_doc = await client.create_document_from_source(
+                assert url_doc.metadata["md5"] == expected_md5
-                "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()
 
 
 @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
+            # First call - should create new document
-    test_content = "Original content for testing."
+            doc1 = await client.create_document_from_source(temp_path)
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+            assert doc1.id is not None
-        f.write(test_content)
+            assert doc1.content == test_content
-        temp_path = Path(f.name)
+            original_id = doc1.id
 
-    try:
+            # Second call with same content - should return existing document (no-op)
-        # First call - should create new document
+            doc2 = await client.create_document_from_source(temp_path)
-        doc1 = await client.create_document_from_source(temp_path)
+            assert doc2.id == original_id  # Same document
-        assert doc1.id is not None
+            assert doc2.content == test_content
-        assert doc1.content == test_content
-        original_id = doc1.id
 
-        # Second call with same content - should return existing document (no-op)
+            # Modify file content
-        doc2 = await client.create_document_from_source(temp_path)
+            updated_content = "Updated content for testing."
-        assert doc2.id == original_id  # Same document
+            temp_path.write_text(updated_content)
-        assert doc2.content == test_content
 
-        # Modify file content
+            # Third call with changed content - should update existing document
-        updated_content = "Updated content for testing."
+            doc3 = await client.create_document_from_source(temp_path)
-        temp_path.write_text(updated_content)
+            assert doc3.id == original_id  # Same document ID
+            assert doc3.content == updated_content  # Updated content
 
-        # Third call with changed content - should update existing document
+            # Verify the document was actually updated in database
-        doc3 = await client.create_document_from_source(temp_path)
+            retrieved_doc = await client.get_document_by_id(original_id)
-        assert doc3.id == original_id  # Same document ID
+            assert retrieved_doc is not None
-        assert doc3.content == updated_content  # Updated content
+            assert retrieved_doc.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()
 
 
 @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"
+        # Mock first response
-    original_content = b"Original URL content"
+        mock_response1 = AsyncMock()
-    updated_content = b"Updated URL content"
+        mock_response1.content = original_content
+        mock_response1.headers = {"content-type": "text/plain"}
+        mock_response1.raise_for_status = AsyncMock()
 
-    # Mock first response
+        with patch("httpx.AsyncClient.get", return_value=mock_response1):
-    mock_response1 = AsyncMock()
+            # First call - should create new document
-    mock_response1.content = original_content
+            doc1 = await client.create_document_from_source(url)
-    mock_response1.headers = {"content-type": "text/plain"}
+            assert doc1.id is not None
-    mock_response1.raise_for_status = AsyncMock()
+            original_id = doc1.id
 
-    with patch("httpx.AsyncClient.get", return_value=mock_response1):
+            # Second call with same content - should return existing document (no-op)
-        # First call - should create new document
+            doc2 = await client.create_document_from_source(url)
-        doc1 = await client.create_document_from_source(url)
+            assert doc2.id == original_id  # Same document
-        assert doc1.id is not None
-        original_id = doc1.id
 
-        # Second call with same content - should return existing document (no-op)
+        mock_response2 = AsyncMock()
-        doc2 = await client.create_document_from_source(url)
+        mock_response2.content = updated_content
-        assert doc2.id == original_id  # Same document
+        mock_response2.headers = {"content-type": "text/plain"}
+        mock_response2.raise_for_status = AsyncMock()
 
-    mock_response2 = AsyncMock()
+        with patch("httpx.AsyncClient.get", return_value=mock_response2):
-    mock_response2.content = updated_content
+            # Third call with changed content - should update existing document
-    mock_response2.headers = {"content-type": "text/plain"}
+            doc3 = await client.create_document_from_source(url)
-    mock_response2.raise_for_status = AsyncMock()
+            assert doc3.id == original_id  # Same document ID
-
+            assert doc3.content == updated_content.decode()  # Updated content
-    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()
 
 
 @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
+        # Create documents
-    doc1_text = "Python is a high-level programming language known for its simplicity and readability."
+        doc1 = await client.create_document(
-    doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
+            content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
-    doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."
+        )
+        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
+        # Test search with keyword that should match doc1
-    doc1 = await client.create_document(
+        results = await client.search("Python programming", limit=3)
-        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
+        assert len(results) > 0
-    results = await client.search("Python programming", limit=3)
+        assert all(len(result) == 2 for result in results)
 
-    assert len(results) > 0
+        # Verify first result is from the Python document (doc1)
-    assert all(len(result) == 2 for result in results)
+        first_chunk, _ = results[0]
+        assert first_chunk.document_id == doc1.id
 
-    # Verify first result is from the Python document (doc1)
+        # Test search with different query
-    first_chunk, _ = results[0]
+        ml_results = await client.search("machine learning data", limit=2)
-    assert first_chunk.document_id == doc1.id
+        assert len(ml_results) > 0
 
-    # Test search with different query
+        # Verify first result is from the machine learning document (doc2)
-    ml_results = await client.search("machine learning data", limit=2)
+        first_ml_chunk, _ = ml_results[0]
-    assert len(ml_results) > 0
+        assert first_ml_chunk.document_id == doc2.id
 
-    # Verify first result is from the machine learning document (doc2)
+        # Test search with limit parameter
-    first_ml_chunk, _ = ml_results[0]
+        limited_results = await client.search("programming", limit=1)
-    assert first_ml_chunk.document_id == doc2.id
+        assert len(limited_results) <= 1
-
-    # Test search with limit parameter
-    limited_results = await client.search("programming", limit=1)
-    assert len(limited_results) <= 1
-
-    client.close()
 
 
 @pytest.mark.asyncio

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:
+    with tempfile.TemporaryDirectory() as temp_dir:
-        f.write("Test content for file watcher")
+        temp_path = Path(temp_dir) / "test.txt"
-        temp_path = Path(f.name)
+        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:
+    with tempfile.TemporaryDirectory() as temp_dir:
-        f.write("Test content for file watcher")
+        temp_path = Path(temp_dir) / "test.txt"
-        temp_path = Path(f.name)
+        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():

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] = []
+        documents_before = await client.list_documents()
-    for content in qa_corpus["document_extracted"][:3]:
+        assert len(documents_before) == 3
-        doc = await client.create_document(
-            content=content,
-        )
-        created_docs.append(doc)
 
-    documents_before = await client.list_documents()
+        chunks_before = []
-    assert len(documents_before) == 3
+        for doc in created_docs:
-
+            assert 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_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:
             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

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
+        try:
-        with pytest.raises(ConfigMismatchError) as exc_info:
+            # Loading the database should raise ConfigMismatchError
-            Store(db_path)
+            with pytest.raises(ConfigMismatchError) as exc_info:
+                Store(db_path)
 
-        assert "CHUNK_SIZE" in str(exc_info.value)
+            assert "CHUNK_SIZE" in str(exc_info.value)
-        assert "Consider rebuilding" in str(exc_info.value)
+            assert "Consider rebuilding" in str(exc_info.value)
 
-        # Rebuild
+            # Rebuild
-        async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
+            async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
-            async for _ in client.rebuild_database():
+                async for _ in client.rebuild_database():
-                pass  # Process all documents
+                    pass  # Process all documents
 
-        # Verify we can now load the database without exception (settings were updated)
+            # Verify we can now load the database without exception (settings were updated)
-        store2 = Store(db_path)
+            store2 = Store(db_path)
-        settings_repo2 = SettingsRepository(store2)
+            settings_repo2 = SettingsRepository(store2)
-        db_settings = settings_repo2.get()
+            db_settings = settings_repo2.get()
-        assert db_settings["CHUNK_SIZE"] == 999
+            assert db_settings["CHUNK_SIZE"] == 999
-        store2.close()
+            store2.close()
 
-        Config.CHUNK_SIZE = original_chunk_size
+        finally:
-
+            Config.CHUNK_SIZE = original_chunk_size
-    finally:
-        # Cleanup
-        if db_path.exists():
-            db_path.unlink()