gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix trailing whitespace and update test mocking for rerank module

Browse source 
• Remove trailing whitespace
• Fix TiktokenTokenizer import patch
• Add async context manager mocks
• Update aiohttp.ClientSession patch
• Improve test reliability
This commit is contained in:
yangdx 2025-12-03 12:40:48 +08:00
parent 8e50eef58b
commit 561ba4e4b5
3 changed files with 23 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
tests/test_rerank_chunking.py
View file

@ -40,7 +40,7 @@ class TestChunkDocumentsForRerank:
long_doc = "a" * 2000 # 2000 characters long_doc = "a" * 2000 # 2000 characters
documents = [long_doc, "short doc"] documents = [long_doc, "short doc"]
with patch("lightrag.rerank.TiktokenTokenizer", side_effect=ImportError): with patch("lightrag.utils.TiktokenTokenizer", side_effect=ImportError):
chunked_docs, doc_indices = chunk_documents_for_rerank( chunked_docs, doc_indices = chunk_documents_for_rerank(
documents, documents,
max_tokens=100, # 100 tokens = ~400 chars max_tokens=100, # 100 tokens = ~400 chars
@ -360,13 +360,17 @@ class TestEndToEndChunking:
mock_response.request_info = None mock_response.request_info = None
mock_response.history = None mock_response.history = None
mock_response.headers = {} mock_response.headers = {}
# Make mock_response an async context manager (for `async with session.post() as response`)
mock_response.__aenter__ = AsyncMock(return_value=mock_response)
mock_response.__aexit__ = AsyncMock(return_value=None)
mock_session = Mock() mock_session = Mock()
mock_session.post = AsyncMock(return_value=mock_response) # session.post() returns an async context manager, so return mock_response which is now one
mock_session.post = Mock(return_value=mock_response)
mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aenter__ = AsyncMock(return_value=mock_session)
mock_session.__aexit__ = AsyncMock() mock_session.__aexit__ = AsyncMock(return_value=None)
with patch("aiohttp.ClientSession", return_value=mock_session): with patch("lightrag.rerank.aiohttp.ClientSession", return_value=mock_session):
result = await cohere_rerank( result = await cohere_rerank(
query=query, query=query,
documents=documents, documents=documents,