From 2053af15a45059b844359bcec19c1ec009077fa7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20MANSUY?= Date: Thu, 4 Dec 2025 19:28:50 +0800 Subject: [PATCH] fix: correct test mock paths for rerank chunking tests --- tests/test_rerank_chunking.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/test_rerank_chunking.py b/tests/test_rerank_chunking.py index e9ad7134..628dde13 100644 --- a/tests/test_rerank_chunking.py +++ b/tests/test_rerank_chunking.py @@ -40,7 +40,7 @@ class TestChunkDocumentsForRerank: long_doc = "a" * 2000 # 2000 characters 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( documents, max_tokens=100, # 100 tokens = ~400 chars @@ -535,10 +535,15 @@ class TestEndToEndChunking: mock_response.history = None mock_response.headers = {} + # Create async context manager for response + mock_response_cm = AsyncMock() + mock_response_cm.__aenter__ = AsyncMock(return_value=mock_response) + mock_response_cm.__aexit__ = AsyncMock(return_value=None) + mock_session = Mock() - mock_session.post = AsyncMock(return_value=mock_response) + mock_session.post = Mock(return_value=mock_response_cm) 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): result = await cohere_rerank(