mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 01:26:48 +00:00
test: isolate chat mocks and cancellation timing
This commit is contained in:
3
tests/fixtures/mock_injector.py
vendored
3
tests/fixtures/mock_injector.py
vendored
@@ -3,7 +3,7 @@ from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from injector import Provider, ScopeDecorator, singleton
|
||||
from injector import Provider, ScopeDecorator, SingletonScope, singleton
|
||||
|
||||
from private_gpt.di import create_application_injector
|
||||
from private_gpt.settings.settings import Settings, unsafe_settings
|
||||
@@ -25,6 +25,7 @@ class MockInjector:
|
||||
if mock is None:
|
||||
mock = MagicMock()
|
||||
self.test_injector.binder.bind(interface, to=mock, scope=scope)
|
||||
self.test_injector.get(SingletonScope)._context.pop(interface, None)
|
||||
return mock # type: ignore
|
||||
|
||||
def bind_settings(self, settings: dict[str, Any]) -> Settings:
|
||||
|
||||
@@ -692,6 +692,7 @@ def test_messages_count_tokens_endpoint_forwards_message_input_and_options(
|
||||
test_client: TestClient,
|
||||
httpx_mock: HTTPXMock,
|
||||
) -> None:
|
||||
injector.get(ChatService)
|
||||
chat_service = Mock()
|
||||
chat_service.count_tokens = AsyncMock(
|
||||
return_value=CountTokensOutput(input_tokens=11)
|
||||
|
||||
@@ -2233,7 +2233,7 @@ async def test_chat_cancels_llm_astream_on_client_disconnection(
|
||||
await request_task
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(llm_generator_closed.wait(), timeout=2.0)
|
||||
await asyncio.wait_for(llm_generator_closed.wait(), timeout=5.0)
|
||||
except TimeoutError:
|
||||
pytest.fail(
|
||||
"LLM astream_chat_with_tools generator was not closed after client disconnection"
|
||||
|
||||
Reference in New Issue
Block a user