5d2e01a6d7
Co-authored-by: Carlo <carloantonio.patti@mistral.ai> Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Clement Sirieix <clem.sirieix@gmail.com> 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: Thomas Kenbeek <thomas.kenbeek@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from contextlib import AbstractContextManager
|
|
import io
|
|
from pathlib import Path
|
|
import re
|
|
import time
|
|
from typing import Protocol
|
|
|
|
import pexpect
|
|
|
|
|
|
class SpawnedVibeProcessFixture(Protocol):
|
|
def __call__(
|
|
self, workdir: Path
|
|
) -> AbstractContextManager[tuple[pexpect.spawn, io.StringIO]]: ...
|
|
|
|
|
|
def ansi_tolerant_pattern(text: str) -> re.Pattern[str]:
|
|
ansi = r"(?:\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*\x07|\r|\n)*"
|
|
return re.compile(ansi.join(re.escape(char) for char in text))
|
|
|
|
|
|
def write_e2e_config(vibe_home: Path, api_base: str) -> None:
|
|
vibe_home.mkdir(parents=True, exist_ok=True)
|
|
(vibe_home / "config.toml").write_text(
|
|
"\n".join([
|
|
'active_model = "mock-model"',
|
|
"enable_update_checks = false",
|
|
"enable_auto_update = false",
|
|
"",
|
|
"[[providers]]",
|
|
'name = "mock-provider"',
|
|
f'api_base = "{api_base}"',
|
|
'api_key_env_var = "MISTRAL_API_KEY"',
|
|
'backend = "generic"',
|
|
"",
|
|
"[[models]]",
|
|
'name = "mock-model"',
|
|
'provider = "mock-provider"',
|
|
'alias = "mock-model"',
|
|
]),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def strip_ansi(text: str) -> str:
|
|
return re.sub(r"\x1b\[[0-9;?]*[ -/]*[@-~]|\x1b\][^\x07]*\x07", "", text)
|
|
|
|
|
|
def wait_for_request_count(
|
|
request_count_getter: Callable[[], int], expected_count: int, timeout: float
|
|
) -> None:
|
|
start = time.monotonic()
|
|
while time.monotonic() - start < timeout:
|
|
if request_count_getter() >= expected_count:
|
|
return
|
|
time.sleep(0.05)
|
|
raise AssertionError(f"Timed out waiting for {expected_count} backend request(s).")
|
|
|
|
|
|
def wait_for_main_screen(child: pexpect.spawn, timeout: float = 20.0) -> None:
|
|
child.expect(ansi_tolerant_pattern("Mistral Vibe v"), timeout=timeout)
|
|
|
|
|
|
def wait_for_rendered_text(
|
|
child: pexpect.spawn, captured: io.StringIO, needle: str, timeout: float
|
|
) -> None:
|
|
start = time.monotonic()
|
|
while time.monotonic() - start < timeout:
|
|
if needle in strip_ansi(captured.getvalue()):
|
|
return
|
|
try:
|
|
child.expect(r"\S", timeout=0.1)
|
|
except pexpect.TIMEOUT:
|
|
pass
|
|
except pexpect.EOF as exc:
|
|
rendered_tail = strip_ansi(captured.getvalue())[-1200:]
|
|
raise AssertionError(
|
|
f"Child exited while waiting for rendered text: {needle!r}\n\nRendered tail:\n{rendered_tail}"
|
|
) from exc
|
|
rendered_tail = strip_ansi(captured.getvalue())[-1200:]
|
|
raise AssertionError(
|
|
f"Timed out waiting for rendered text: {needle!r}\n\nRendered tail:\n{rendered_tail}"
|
|
)
|