github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-17 01:48:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: tests

Browse Source
This commit is contained in:
Javier Martinez
2026-07-13 11:33:04 +02:00
parent 0e4f26eaac
commit 5a620ec485
9 changed files with 202 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
private_gpt/components/ingestion/ingestion_scheduler.py
View File

@@ -29,10 +29,19 @@ class BaseIngestionScheduler(ABC):
The async methods always dispatch to Celery regardless of mode.
"""
def __init__(self, ingest_service: IngestService, s3_helper: S3Helper) -> None:
def __init__(
self,
ingest_service: IngestService,
s3_helper: S3Helper | None = None,
) -> None:
self._ingest_service = ingest_service
self._s3_helper = s3_helper
def _require_s3_helper(self) -> S3Helper:
if self._s3_helper is None:
raise RuntimeError("S3Helper is required for remote ingestion flows")
return self._s3_helper
@abstractmethod
def ingest(self, ingest_body: IngestBody) -> IngestResponse:
...
@@ -61,7 +70,7 @@ class BaseIngestionScheduler(ABC):
if should_upload:
content = ingest_body.ingest_body.input.to_binary_content()
object_name = str(uuid.uuid4())
s3_url = self._s3_helper.upload_file_to_s3(
s3_url = self._require_s3_helper().upload_file_to_s3(
filename=content.filename,
bytes_data=content.data.read(),
bucket_name=config.s3.temporary_bucket_name,
@@ -118,8 +127,8 @@ class LocalIngestionScheduler(BaseIngestionScheduler):
"""Run sync ingestion in-process."""
@inject
def __init__(self, ingest_service: IngestService, s3_helper: S3Helper) -> None:
super().__init__(ingest_service, s3_helper)
def __init__(self, ingest_service: IngestService) -> None:
super().__init__(ingest_service)
def ingest(self, ingest_body: IngestBody) -> IngestResponse:
from private_gpt.server.ingest.ingest_router import ingest_data_sync
@@ -189,7 +198,7 @@ class CeleryIngestionScheduler(BaseIngestionScheduler):
else None
)
object_name = str(uuid.uuid4())
s3_url = self._s3_helper.upload_file_to_s3(
s3_url = self._require_s3_helper().upload_file_to_s3(
filename=content.filename,
bytes_data=content.data.read(),
bucket_name=config.s3.temporary_bucket_name,

6
private_gpt/components/storage/s3_helper.py
View File

@@ -32,9 +32,9 @@ class S3Helper:
def _get_s3_client(s3_settings: S3Settings) -> S3Client | None:
if (
not s3_settings
or s3_settings.access_key_id is None
or s3_settings.secret_access_key is None
or s3_settings.endpoint_url is None
or not s3_settings.access_key_id
or not s3_settings.secret_access_key
or not s3_settings.endpoint_url
):
return None

223
tests/components/streaming/test_chat_scheduler.py
View File

@@ -1,71 +1,14 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock
import pytest
from private_gpt.chat.input_models import MessageInput
from private_gpt.components.streaming.providers.models import StreamStatus
from private_gpt.components.streaming.tasks.chat_scheduler import (
ArqChatScheduler,
ChatSchedulerFactory,
LocalChatScheduler,
)
from private_gpt.server.chat.chat_models import ChatBody
def _chat_body() -> ChatBody:
return ChatBody(messages=[MessageInput(role="user", content="hello")])
def _make_stream_component(correlation_id: str = "msg-1") -> MagicMock:
stream_service = AsyncMock()
stream_service.create_stream = AsyncMock(return_value=correlation_id)
stream_service.update_stream_status = AsyncMock()
stream_component = MagicMock()
stream_component.stream = stream_service
return stream_component
@pytest.mark.anyio
async def test_arq_chat_scheduler_create_enqueues_job(
monkeypatch: pytest.MonkeyPatch,
) -> None:
enqueue_chat_job = AsyncMock()
monkeypatch.setattr(
"private_gpt.components.streaming.tasks.chat_scheduler.enqueue_chat_job",
enqueue_chat_job,
)
stream_component = _make_stream_component("msg-arq-1")
scheduler = ArqChatScheduler(stream_component=stream_component)
correlation_id = await scheduler.create(_chat_body(), message_id="msg-arq-1")
assert correlation_id == "msg-arq-1"
enqueue_chat_job.assert_awaited_once()
@pytest.mark.anyio
async def test_arq_chat_scheduler_create_marks_error_on_enqueue_failure(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(
"private_gpt.components.streaming.tasks.chat_scheduler.enqueue_chat_job",
AsyncMock(side_effect=RuntimeError("redis unavailable")),
)
stream_component = _make_stream_component("msg-arq-2")
scheduler = ArqChatScheduler(stream_component=stream_component)
with pytest.raises(RuntimeError, match="redis unavailable"):
await scheduler.create(_chat_body(), message_id="msg-arq-2")
stream_component.stream.update_stream_status.assert_awaited_once_with(
"msg-arq-2",
StreamStatus.ERROR,
error_message="redis unavailable",
metadata={"message_count": 1},
)
@pytest.mark.anyio
@@ -78,96 +21,114 @@ async def test_arq_chat_scheduler_cancel_aborts_job(
abort_chat_job,
)
scheduler = ArqChatScheduler(stream_component=_make_stream_component())
scheduler = ArqChatScheduler()
cancelled = await scheduler.cancel("msg-arq-3")
assert cancelled is True
abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-3")
@pytest.mark.anyio
async def test_local_chat_scheduler_create_starts_task() -> None:
chat_service = AsyncMock()
chat_request_mapper = AsyncMock()
chat_request_mapper.create_request_from_body.return_value = MagicMock()
start_gen = AsyncMock()
start_state = MagicMock()
start_state.runtime.iteration = 0
start_state.runtime.next_block_count = 0
start_state.input.request.model_dump.return_value = {}
start_gen.final_state_task = AsyncMock(return_value=start_state)()
chat_service.start_chat.return_value = start_gen
iter_gen = AsyncMock()
iter_state = MagicMock()
from private_gpt.components.engines.chat.models.chat_state import (
ChatStatus,
)
iter_state.output.status = ChatStatus.COMPLETED
iter_state.output.stop_reason = "end_turn"
iter_state.runtime.iteration = 1
iter_state.runtime.next_block_count = 0
iter_state.input.request.model_dump.return_value = {}
iter_gen.final_state_task = AsyncMock(return_value=iter_state)()
chat_service.run_iteration.return_value = iter_gen
close_gen = AsyncMock()
close_state = MagicMock()
close_gen.final_state_task = AsyncMock(return_value=close_state)()
chat_service.close_chat.return_value = close_gen
stream_component = _make_stream_component("msg-local-1")
stream_processor = MagicMock()
stream_processor.process_stream = AsyncMock()
stream_processor.stream_service = stream_component.stream
task_manager = AsyncMock()
task_manager.create_task = AsyncMock()
iteration_state_service = AsyncMock()
scheduler = LocalChatScheduler(
chat_service=chat_service,
chat_request_mapper=chat_request_mapper,
stream_component=stream_component,
stream_processor=stream_processor,
task_manager=task_manager,
iteration_state_service=iteration_state_service,
)
correlation_id = await scheduler.create(_chat_body(), message_id="msg-local-1")
assert correlation_id == "msg-local-1"
task_manager.create_task.assert_awaited_once()
@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_cancels_task() -> None:
task_manager = AsyncMock()
task_manager.cancel_task = AsyncMock(return_value=True)
async def _work() -> None:
await asyncio.sleep(100)
scheduler = LocalChatScheduler(
chat_service=AsyncMock(),
chat_request_mapper=AsyncMock(),
stream_component=_make_stream_component(),
stream_processor=MagicMock(),
task_manager=task_manager,
iteration_state_service=AsyncMock(),
)
task = asyncio.create_task(_work(), name="chat_msg-local-2")
scheduler = LocalChatScheduler()
cancelled = await scheduler.cancel("msg-local-2")
assert cancelled is True
task_manager.cancel_task.assert_awaited_once_with("msg-local-2")
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
def test_chat_scheduler_factory_selects_mode() -> None:
@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_returns_false_when_no_task() -> None:
scheduler = LocalChatScheduler()
cancelled = await scheduler.cancel("nonexistent-msg")
assert cancelled is False
def test_chat_scheduler_factory_selects_local_mode() -> None:
factory = ChatSchedulerFactory(
settings=SimpleNamespace(
scheduler=SimpleNamespace(chat=SimpleNamespace(mode="local"))
),
)
assert isinstance(factory.get(), LocalChatScheduler)
def test_chat_scheduler_factory_selects_arq_mode() -> None:
factory = ChatSchedulerFactory(
settings=SimpleNamespace(
scheduler=SimpleNamespace(chat=SimpleNamespace(mode="arq"))
),
local=MagicMock(),
arq=ArqChatScheduler(stream_component=_make_stream_component()),
)
assert isinstance(factory.get(), ArqChatScheduler)
def test_chat_scheduler_factory_raises_on_unknown_mode() -> None:
with pytest.raises(ValueError, match="Unknown"):
ChatSchedulerFactory(
settings=SimpleNamespace(
scheduler=SimpleNamespace(chat=SimpleNamespace(mode="unknown"))
),
)
@pytest.mark.anyio
async def test_arq_chat_scheduler_cancel_returns_false_when_abort_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
abort_chat_job = AsyncMock(return_value=False)
monkeypatch.setattr(
"private_gpt.components.streaming.tasks.chat_scheduler.abort_chat_job",
abort_chat_job,
)
assert isinstance(factory.get(), ArqChatScheduler)
scheduler = ArqChatScheduler()
cancelled = await scheduler.cancel("msg-arq-not-found")
assert cancelled is False
abort_chat_job.assert_awaited_once_with(correlation_id="msg-arq-not-found")
@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_marks_task_as_cancelled() -> None:
async def _work() -> None:
await asyncio.sleep(100)
task = asyncio.create_task(_work(), name="chat_msg-local-mark")
scheduler = LocalChatScheduler()
result = await scheduler.cancel("msg-local-mark")
assert result is True
with pytest.raises(asyncio.CancelledError):
await task
assert task.cancelled()
@pytest.mark.anyio
async def test_local_chat_scheduler_cancel_only_targets_the_matching_task() -> None:
async def _work() -> None:
await asyncio.sleep(100)
task_target = asyncio.create_task(_work(), name="chat_target-cid")
task_other = asyncio.create_task(_work(), name="chat_other-cid")
scheduler = LocalChatScheduler()
result = await scheduler.cancel("target-cid")
assert result is True
with pytest.raises(asyncio.CancelledError):
await task_target
assert task_target.cancelled()
assert not task_other.done()
task_other.cancel()
with pytest.raises(asyncio.CancelledError):
await task_other

63
tests/components/tools/test_tool_scheduler.py
View File

@@ -1,6 +1,5 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import MagicMock
import pytest
@@ -9,39 +8,12 @@ from private_gpt.components.tools.remote_execution import ToolExecutionRequest
from private_gpt.components.tools.tool_scheduler import CeleryToolScheduler
class _Result:
def __init__(self, task_id: str) -> None:
self.id = task_id
self.result = None
def ready(self) -> bool:
return False
def failed(self) -> bool:
return False
@pytest.mark.anyio
async def test_celery_tool_scheduler_cancels_child_task_when_stream_cancelled(
monkeypatch: pytest.MonkeyPatch,
) -> None:
dispatch_task = MagicMock(return_value=_Result("msg-1:tool-1"))
monkeypatch.setattr(
"private_gpt.components.tools.tool_scheduler.dispatch_task",
dispatch_task,
)
celery_app = MagicMock()
monkeypatch.setattr("private_gpt.celery.celery.celery_app", celery_app)
stream_component = SimpleNamespace(
stream=SimpleNamespace(is_cancelled=AsyncMock(return_value=True))
)
async def test_celery_tool_scheduler_execute_raises_not_implemented() -> None:
scheduler = CeleryToolScheduler(
settings=SimpleNamespace(
scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools"))
),
stream_component=stream_component,
)
request = ToolExecutionRequest.model_validate(
@@ -54,7 +26,34 @@ async def test_celery_tool_scheduler_cancels_child_task_when_stream_cancelled(
}
)
with pytest.raises(asyncio.CancelledError):
with pytest.raises(NotImplementedError):
await scheduler.execute(request)
celery_app.control.revoke.assert_called_once_with("msg-1:tool-1", terminate=True)
@pytest.mark.anyio
async def test_celery_tool_scheduler_cancel_revokes_task(
monkeypatch: pytest.MonkeyPatch,
) -> None:
celery_app = MagicMock()
monkeypatch.setattr("private_gpt.celery.celery.celery_app", celery_app)
scheduler = CeleryToolScheduler(
settings=SimpleNamespace(
scheduler=SimpleNamespace(tools=SimpleNamespace(celery_queue="tools"))
),
)
request = ToolExecutionRequest.model_validate(
{
"tool_id": "tool-1",
"tool_name": "bash",
"tool_kwargs": {},
"tool_spec": ToolSpec(name="bash", runtime="server", input_schema={}),
"context": {"correlation_id": "msg-1", "messages": []},
}
)
cancelled = await scheduler.cancel(request, task_id="task-abc")
assert cancelled is True
celery_app.control.revoke.assert_called_once_with("task-abc", terminate=True)

96
tests/server/chat/test_chat_async_worker.py
View File

@@ -1,93 +1,87 @@
from unittest.mock import AsyncMock
from collections.abc import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
from private_gpt.chat.input_models import MessageInput
from private_gpt.components.streaming.providers.in_memory_stream_service import (
InMemoryStreamService,
)
from private_gpt.server.chat.chat_models import ChatBody
from private_gpt.events.models import Event
from private_gpt.server.chat_async.chat_async_service import ChatAsyncService
class _StreamManagerStub:
def __init__(self, stream_service: InMemoryStreamService) -> None:
self.stream_service = stream_service
self.cancel_stream = AsyncMock(return_value=False)
self.stream_exists = AsyncMock(return_value=False)
class _ChatSchedulerFactoryStub:
def __init__(self, scheduler: object) -> None:
self._scheduler = scheduler
def get(self) -> object:
return self._scheduler
class _ChatSchedulerStub:
def __init__(self) -> None:
self.create = AsyncMock(return_value="msg-1")
async def _empty_gen() -> AsyncGenerator[Event, None]:
return
yield # make it an async generator
def _service(
stream_service: InMemoryStreamService,
scheduler: _ChatSchedulerStub | None = None,
stream_manager: MagicMock | None = None,
chat_facade: MagicMock | None = None,
) -> ChatAsyncService:
scheduler = scheduler or _ChatSchedulerStub()
if stream_manager is None:
stream_manager = MagicMock()
stream_manager.stream_exists = AsyncMock(return_value=False)
stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1")
stream_manager.cancel_stream = AsyncMock(return_value=False)
if chat_facade is None:
chat_facade = MagicMock()
chat_facade.create_chat_event_generator = AsyncMock(return_value=_empty_gen())
return ChatAsyncService(
stream_manager=_StreamManagerStub(stream_service), # type: ignore[arg-type]
chat_scheduler_factory=_ChatSchedulerFactoryStub(scheduler), # type: ignore[arg-type]
chat_facade=chat_facade,
stream_manager=stream_manager,
)
def _chat_body() -> ChatBody:
return ChatBody(messages=[MessageInput(role="user", content="hello")])
@pytest.mark.anyio
async def test_initiate_chat_stream_delegates_to_scheduler() -> None:
stream_service = InMemoryStreamService()
scheduler = _ChatSchedulerStub()
service = _service(stream_service, scheduler)
stream_manager = MagicMock()
stream_manager.stream_exists = AsyncMock(return_value=False)
stream_manager.create_and_start_stream = AsyncMock(return_value="msg-1")
chat_facade = MagicMock()
chat_facade.create_chat_event_generator = AsyncMock(return_value=_empty_gen())
service = _service(stream_manager=stream_manager, chat_facade=chat_facade)
correlation_id = await service.initiate_chat_stream(
body=_chat_body(),
request=MagicMock(),
message_id="msg-1",
)
assert correlation_id == "msg-1"
scheduler.create.assert_awaited_once_with(body=_chat_body(), message_id="msg-1")
stream_manager.create_and_start_stream.assert_awaited_once()
@pytest.mark.anyio
async def test_initiate_chat_stream_rejects_duplicate_message_id() -> None:
stream_service = InMemoryStreamService()
scheduler = _ChatSchedulerStub()
service = _service(stream_service, scheduler)
service.stream_manager.stream_exists.return_value = True
stream_manager = MagicMock()
stream_manager.stream_exists = AsyncMock(return_value=True)
service = _service(stream_manager=stream_manager)
with pytest.raises(ValueError, match="already exists"):
await service.initiate_chat_stream(body=_chat_body(), message_id="msg-2")
scheduler.create.assert_not_awaited()
await service.initiate_chat_stream(request=MagicMock(), message_id="msg-2")
@pytest.mark.anyio
async def test_worker_cancel_marks_stream_cancelled() -> None:
stream_service = InMemoryStreamService()
service = _service(stream_service)
stream_manager = MagicMock()
stream_manager.cancel_stream = AsyncMock(return_value=False)
service = _service(stream_manager=stream_manager)
await service.cancel_stream("msg-3")
service.stream_manager.cancel_stream.assert_awaited_once_with("msg-3")
stream_manager.cancel_stream.assert_awaited_once_with("msg-3")
@pytest.mark.anyio
async def test_arq_cancel_delegates_to_stream_manager() -> None:
stream_service = InMemoryStreamService()
service = _service(stream_service)
stream_manager = MagicMock()
stream_manager.cancel_stream = AsyncMock(return_value=False)
service = _service(stream_manager=stream_manager)
await service.cancel_stream("msg-4")
service.stream_manager.cancel_stream.assert_awaited_once_with("msg-4")
stream_manager.cancel_stream.assert_awaited_once_with("msg-4")

16
tests/server/chat/test_chat_service.py
View File

@@ -1,5 +1,5 @@
from collections.abc import AsyncGenerator
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock
import pytest
from llama_index.core.base.llms.types import MessageRole
@@ -9,6 +9,7 @@ from private_gpt.components.chat.models.chat_config_models import (
ResolvedChatRequest,
ResolvedSystemConfig,
)
from private_gpt.components.engines.chat.async_chat_engine import ChatExecution
from private_gpt.components.engines.chat.interceptors.chat_interceptor import (
ChatRequestLoopInterceptor,
)
@@ -66,7 +67,8 @@ async def _fatal_event_stream() -> AsyncGenerator[Event, None]:
def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock:
mock_engine = MagicMock()
mock_engine.run.return_value = stream
execution = ChatExecution(events=stream, final_state_task=MagicMock())
mock_engine.run = AsyncMock(return_value=execution)
return mock_engine
@@ -74,7 +76,7 @@ def _mock_engine_for(stream: AsyncGenerator[Event, None]) -> MagicMock:
async def test_chat_folds_loop_events(injector: MockInjector) -> None:
service: ChatService = injector.get(ChatService)
mock_engine = _mock_engine_for(_basic_event_stream())
service._build_engine = MagicMock(return_value=mock_engine)
service.build_engine = MagicMock(return_value=mock_engine)
request = ResolvedChatRequest(
messages=[ChatMessage(content="hi", role=MessageRole.USER)],
@@ -87,8 +89,8 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None:
assert result.response == "hello"
assert result.stop_reason == "end_turn"
assert result.usage is not None
assert result.usage["input_tokens"] == 10
assert result.usage["output_tokens"] == 20
assert result.usage.input_tokens == 10
assert result.usage.output_tokens == 20
mock_engine.run.assert_called_once()
@@ -96,7 +98,7 @@ async def test_chat_folds_loop_events(injector: MockInjector) -> None:
async def test_stream_chat_returns_loop_generator(injector: MockInjector) -> None:
service: ChatService = injector.get(ChatService)
mock_engine = _mock_engine_for(_basic_event_stream())
service._build_engine = MagicMock(return_value=mock_engine)
service.build_engine = MagicMock(return_value=mock_engine)
request = ResolvedChatRequest(
messages=[ChatMessage(content="hi", role=MessageRole.USER)],
@@ -138,7 +140,7 @@ async def test_chat_propagates_fatal_error_to_completion(
) -> None:
service: ChatService = injector.get(ChatService)
mock_engine = _mock_engine_for(_fatal_event_stream())
service._build_engine = MagicMock(return_value=mock_engine)
service.build_engine = MagicMock(return_value=mock_engine)
request = ResolvedChatRequest(
messages=[ChatMessage(content="hi", role=MessageRole.USER)],

7
tests/settings/test_settings.py
View File

@@ -31,17 +31,16 @@ def test_chat_scheduler_requires_redis_stream_backend() -> None:
Settings(**merged)
def test_chat_scheduler_requires_real_celery_worker() -> None:
def test_chat_scheduler_celery_mode_requires_redis_broker() -> None:
merged = merge_settings(
[
unsafe_settings,
{
"scheduler": {"chat": {"mode": "celery"}},
"stream": {"broker": "redis"},
"celery": {"use_workers": False},
"stream": {"broker": "memory"},
},
]
)
with pytest.raises(ValueError, match=r"celery\.use_workers=true"):
with pytest.raises(ValueError, match=r"stream\.broker=redis"):
Settings(**merged)

3
tests/sse/test_sse.py
View File

@@ -163,7 +163,8 @@ async def _noop_tool(value: str) -> str:
async def _collect_engine(engine: ChatLoopEngine, request: ResolvedChatRequest) -> list:
result = []
async for event in engine.run(request):
execution = await engine.run(request)
async for event in execution.events:
result.append(event)
return result

10
tests/utils/test_tokens.py
View File

@@ -3,7 +3,7 @@ from unittest.mock import MagicMock
import pytest
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from private_gpt.components.llm.llm_helper import get_tokenizer_fn
from private_gpt.components.llm.llm_helper import get_async_tokenizer_fn
from private_gpt.components.llm.tokenizers.tokenizer_base import (
TokenizedInput,
TokenizerBase,
@@ -32,7 +32,7 @@ class AsyncCapableTokenizer(MagicMock):
@pytest.mark.asyncio
async def test_async_tokenizer_prefers_underlying_acall():
tokenizer = AsyncCapableTokenizer()
tokenizer_fn = get_tokenizer_fn(tokenizer)
tokenizer_fn = get_async_tokenizer_fn(tokenizer)
tokens = await async_tokenizer("one two three", tokenizer_fn=tokenizer_fn)
@@ -44,11 +44,7 @@ async def test_async_tokenizer_prefers_underlying_acall():
@pytest.mark.asyncio
async def test_estimate_token_count_uses_async_tokenizer_wrapper():
tokenizer = AsyncCapableTokenizer()
tokenizer_fn = get_tokenizer_fn(tokenizer)
async def message_to_input(messages, **kwargs):
del kwargs
raise AssertionError("message_to_input should run synchronously in this test")
tokenizer_fn = get_async_tokenizer_fn(tokenizer)
count = await estimate_token_count(
chat_history=[ChatMessage(role=MessageRole.USER, content="one two three")],