Initial commit

Co-Authored-By: Quentin Torroba <quentin.torroba@mistral.ai>
Co-Authored-By: Laure Hugo <laure.hugo@mistral.ai>
Co-Authored-By: Benjamin Trom <benjamin.trom@mistral.ai>
Co-Authored-By: Mathias Gesbert <mathias.gesbert@ext.mistral.ai>
Co-Authored-By: Michel Thomazo <michel.thomazo@mistral.ai>
Co-Authored-By: Clément Drouin <clement.drouin@mistral.ai>
Co-Authored-By: Vincent Guilloux <vincent.guilloux@mistral.ai>
Co-Authored-By: Valentin Berard <val@mistral.ai>
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>

tests/mock/mock_backend_factory.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from vibe.core.config import Backend
+from vibe.core.llm.backend.factory import BACKEND_FACTORY
+
+
+@contextmanager
+def mock_backend_factory(backend_type: Backend, factory_func):
+    original = BACKEND_FACTORY[backend_type]
+    try:
+        BACKEND_FACTORY[backend_type] = factory_func
+        yield
+    finally:
+        BACKEND_FACTORY[backend_type] = original

tests/mock/mock_entrypoint.py

@@ -0,0 +1,66 @@
+"""Wrapper script that intercepts LLM calls when mocking is enabled.
+
+This script is used to mock the LLM calls when testing the CLI.
+Mocked returns are stored in the VIBE_MOCK_LLM_DATA environment variable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+import json
+import os
+import sys
+from unittest.mock import patch
+
+from pydantic import ValidationError
+
+from tests import TESTS_ROOT
+from tests.mock.utils import MOCK_DATA_ENV_VAR
+from vibe.core.types import LLMChunk
+
+
+def mock_llm_output() -> None:
+    sys.path.insert(0, str(TESTS_ROOT))
+
+    # Apply mocking before importing any vibe modules
+    mock_data_str = os.environ.get(MOCK_DATA_ENV_VAR)
+    if not mock_data_str:
+        raise ValueError(f"{MOCK_DATA_ENV_VAR} is not set")
+    mock_data = json.loads(mock_data_str)
+    try:
+        chunks = [LLMChunk.model_validate(chunk) for chunk in mock_data]
+    except ValidationError as e:
+        raise ValueError(f"Invalid mock data: {e}") from e
+
+    chunk_iterable = iter(chunks)
+
+    async def mock_complete(*args, **kwargs) -> LLMChunk:
+        return next(chunk_iterable)
+
+    async def mock_complete_streaming(*args, **kwargs) -> AsyncGenerator[LLMChunk]:
+        yield next(chunk_iterable)
+
+    patch(
+        "vibe.core.llm.backend.mistral.MistralBackend.complete",
+        side_effect=mock_complete,
+    ).start()
+    patch(
+        "vibe.core.llm.backend.generic.GenericBackend.complete",
+        side_effect=mock_complete,
+    ).start()
+    patch(
+        "vibe.core.llm.backend.mistral.MistralBackend.complete_streaming",
+        side_effect=mock_complete_streaming,
+    ).start()
+    patch(
+        "vibe.core.llm.backend.generic.GenericBackend.complete_streaming",
+        side_effect=mock_complete_streaming,
+    ).start()
+
+
+if __name__ == "__main__":
+    mock_llm_output()
+
+    from vibe.acp.entrypoint import main
+
+    main()

tests/mock/utils.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import json
+
+from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role, ToolCall
+
+MOCK_DATA_ENV_VAR = "VIBE_MOCK_LLM_DATA"
+
+
+def mock_llm_chunk(
+    content: str = "Hello!",
+    role: Role = Role.assistant,
+    tool_calls: list[ToolCall] | None = None,
+    name: str | None = None,
+    tool_call_id: str | None = None,
+    finish_reason: str | None = None,
+    prompt_tokens: int = 10,
+    completion_tokens: int = 5,
+) -> LLMChunk:
+    message = LLMMessage(
+        role=role,
+        content=content,
+        tool_calls=tool_calls,
+        name=name,
+        tool_call_id=tool_call_id,
+    )
+    return LLMChunk(
+        message=message,
+        usage=LLMUsage(
+            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
+        ),
+        finish_reason=finish_reason,
+    )
+
+
+def get_mocking_env(mock_chunks: list[LLMChunk] | None = None) -> dict[str, str]:
+    if mock_chunks is None:
+        mock_chunks = [mock_llm_chunk()]
+
+    mock_data = [LLMChunk.model_dump(mock_chunk) for mock_chunk in mock_chunks]
+
+    return {MOCK_DATA_ENV_VAR: json.dumps(mock_data)}