danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix client test on context expansion to not depend on .env settings

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-12 12:38:22 +03:00
parent d90dfe258c
commit c43ae62018
No known key found for this signature in database
1 changed files with 85 additions and 89 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/test_client.py
View file

@ -526,8 +526,6 @@ async def test_client_ask_with_cite(temp_db_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context(temp_db_path): async def test_client_expand_context(temp_db_path):
"""Test expanding search results with adjacent chunks.""" """Test expanding search results with adjacent chunks."""
# Mock Config to have CONTEXT_CHUNK_RADIUS = 2
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2):
async with HaikuRAG(temp_db_path) as client: async with HaikuRAG(temp_db_path) as client:
# Create chunks manually # Create chunks manually
manual_chunks = [ manual_chunks = [
@ -553,8 +551,8 @@ async def test_client_expand_context(temp_db_path):
middle_chunk = next(c for c in chunks if c.metadata.get("order") == 2) middle_chunk = next(c for c in chunks if c.metadata.get("order") == 2)
search_results = [(middle_chunk, 0.8)] search_results = [(middle_chunk, 0.8)]
# Test expand_context # Test expand_context with radius=2
expanded_results = await client.expand_context(search_results) expanded_results = await client.expand_context(search_results, radius=2)
assert len(expanded_results) == 1 assert len(expanded_results) == 1
expanded_chunk, score = expanded_results[0] expanded_chunk, score = expanded_results[0]
@ -575,7 +573,6 @@ async def test_client_expand_context(temp_db_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_radius_zero(temp_db_path): async def test_client_expand_context_radius_zero(temp_db_path):
"""Test expand_context with radius 0 returns original results.""" """Test expand_context with radius 0 returns original results."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 0):
async with HaikuRAG(temp_db_path) as client: async with HaikuRAG(temp_db_path) as client:
# Create a simple document # Create a simple document
doc = await client.create_document(content="Simple test content") doc = await client.create_document(content="Simple test content")
@ -583,7 +580,7 @@ async def test_client_expand_context_radius_zero(temp_db_path):
chunks = await client.chunk_repository.get_by_document_id(doc.id) chunks = await client.chunk_repository.get_by_document_id(doc.id)
search_results = [(chunks[0], 0.9)] search_results = [(chunks[0], 0.9)]
expanded_results = await client.expand_context(search_results) expanded_results = await client.expand_context(search_results, radius=0)
# Should return exactly the same results # Should return exactly the same results
assert expanded_results == search_results assert expanded_results == search_results
@ -592,7 +589,6 @@ async def test_client_expand_context_radius_zero(temp_db_path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_client_expand_context_multiple_chunks(temp_db_path): async def test_client_expand_context_multiple_chunks(temp_db_path):
"""Test expand_context with multiple search results.""" """Test expand_context with multiple search results."""
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1):
async with HaikuRAG(temp_db_path) as client: async with HaikuRAG(temp_db_path) as client:
# Create first document with manual chunks # Create first document with manual chunks
doc1_chunks = [ doc1_chunks = [
@ -623,7 +619,7 @@ async def test_client_expand_context_multiple_chunks(temp_db_path):
chunk2 = next(c for c in chunks2 if c.metadata.get("order") == 0) chunk2 = next(c for c in chunks2 if c.metadata.get("order") == 0)
search_results = [(chunk1, 0.8), (chunk2, 0.7)] search_results = [(chunk1, 0.8), (chunk2, 0.7)]
expanded_results = await client.expand_context(search_results) expanded_results = await client.expand_context(search_results, radius=1)
assert len(expanded_results) == 2 assert len(expanded_results) == 2