vibe/tests/core/test_teleport_service.py
Mathias Gesbert cf3f4ca58f
v2.11.1 (#721)
Co-authored-by: Clément Drouin <clement.drouin@mistral.ai>
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com>
Co-authored-by: Mert Unsal <mertunsal1905@gmail.com>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Val <102326092+vdeva@users.noreply.github.com>
Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com>
Co-authored-by: maximevoisin-pm <maxime.voisin@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
2026-05-27 17:10:40 +02:00

725 lines
27 KiB
Python

from __future__ import annotations
import base64
import importlib
import json
import os
from pathlib import Path
import sys
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import zstandard
from vibe.core.config import VibeConfig
from vibe.core.teleport.errors import (
ServiceTeleportError,
ServiceTeleportNotSupportedError,
)
from vibe.core.teleport.git import GitRepoInfo
from vibe.core.teleport.nuage import GitHubStatus, TeleportSession
from vibe.core.teleport.teleport import TeleportService
from vibe.core.teleport.types import (
TeleportAuthCompleteEvent,
TeleportAuthRequiredEvent,
TeleportCheckingGitEvent,
TeleportCompleteEvent,
TeleportFetchingUrlEvent,
TeleportPushingEvent,
TeleportPushRequiredEvent,
TeleportPushResponseEvent,
TeleportStartingWorkflowEvent,
TeleportWaitingForGitHubEvent,
)
def _reimport_agent_loop() -> Any:
to_clear = ("vibe.core.agent_loop", "git", "vibe.core.teleport")
for k in [k for k in sys.modules if any(k.startswith(m) for m in to_clear)]:
del sys.modules[k]
return importlib.import_module("vibe.core.agent_loop")
class TestTeleportServiceCompressDiff:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
def test_returns_none_for_empty_diff(self, service: TeleportService) -> None:
assert service._compress_diff("") is None
def test_compresses_and_encodes_diff(self, service: TeleportService) -> None:
diff = "diff --git a/file.txt b/file.txt\n+new line"
result = service._compress_diff(diff)
assert result is not None
decoded = base64.b64decode(result)
decompressed = zstandard.ZstdDecompressor().decompress(decoded)
assert decompressed.decode("utf-8") == diff
def test_raises_when_diff_too_large(self, service: TeleportService) -> None:
large_diff = "x" * 2_000_000
with pytest.raises(ServiceTeleportError, match="Diff too large"):
service._compress_diff(large_diff, max_size=100)
class TestTeleportServiceBuildGitHubParams:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
def test_builds_params_from_git_info(self, service: TeleportService) -> None:
git_info = GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
params = service._build_github_params(git_info)
assert params.repo == "owner/repo"
assert params.branch == "main"
assert params.commit == "abc123"
assert params.teleported_diffs is None
def test_includes_compressed_diff(self, service: TeleportService) -> None:
git_info = GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="diff content",
)
params = service._build_github_params(git_info)
assert params.teleported_diffs is not None
class TestTeleportServiceValidateConfig:
def test_raises_when_no_api_key(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="",
workdir=tmp_path,
)
with pytest.raises(ServiceTeleportError, match="MISTRAL_API_KEY not set"):
service._validate_config()
def test_passes_when_api_key_set(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="valid-key",
workdir=tmp_path,
)
service._validate_config()
def test_uses_custom_env_var_name_in_error(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
mock_config = MagicMock()
mock_config.vibe_code_api_key_env_var = "CUSTOM_API_KEY"
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="",
workdir=tmp_path,
vibe_config=mock_config,
)
with pytest.raises(ServiceTeleportError, match="CUSTOM_API_KEY not set"):
service._validate_config()
class TestTeleportServiceCheckSupported:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
@pytest.mark.asyncio
async def test_check_supported_calls_git_info(
self, service: TeleportService
) -> None:
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
)
await service.check_supported()
service._git.get_info.assert_called_once()
@pytest.mark.asyncio
async def test_check_supported_raises_when_not_supported(
self, service: TeleportService
) -> None:
service._git.get_info = AsyncMock(
side_effect=ServiceTeleportNotSupportedError(
"Teleport requires a git repository. cd into a project with a .git directory and try again."
)
)
with pytest.raises(ServiceTeleportNotSupportedError):
await service.check_supported()
class TestTeleportServiceIsSupported:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
return TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
@pytest.mark.asyncio
async def test_is_supported_returns_true(self, service: TeleportService) -> None:
service._git.is_supported = AsyncMock(return_value=True)
assert await service.is_supported() is True
@pytest.mark.asyncio
async def test_is_supported_returns_false(self, service: TeleportService) -> None:
service._git.is_supported = AsyncMock(return_value=False)
assert await service.is_supported() is False
class TestTeleportServiceExecute:
@pytest.fixture
def service(self, tmp_path: Path) -> TeleportService:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
service._git.fetch = AsyncMock()
service._git.is_branch_pushed = AsyncMock(return_value=True)
return service
@pytest.fixture
def git_info(self) -> GitRepoInfo:
return GitRepoInfo(
remote_url="https://github.com/owner/repo.git",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
@pytest.fixture
def mock_github_connected(self) -> MagicMock:
github_data = MagicMock()
github_data.connected = True
github_data.oauth_url = None
github_data.status = GitHubStatus.CONNECTED
return github_data
@pytest.mark.asyncio
async def test_execute_happy_path_github_already_connected(
self,
service: TeleportService,
git_info: GitRepoInfo,
mock_github_connected: MagicMock,
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=True)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
async for event in gen:
events.append(event)
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
assert isinstance(events[3], TeleportAuthCompleteEvent)
assert isinstance(events[4], TeleportFetchingUrlEvent)
assert isinstance(events[5], TeleportCompleteEvent)
assert events[5].url == "https://chat.example.com/123"
workflow_params = mock_nuage.start_workflow.call_args.args[0]
assert workflow_params.integrations.chat_assistant is not None
assert workflow_params.integrations.chat_assistant.project_name is None
@pytest.mark.asyncio
async def test_execute_requires_push_and_user_approves(
self,
service: TeleportService,
git_info: GitRepoInfo,
mock_github_connected: MagicMock,
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=False)
service._git.get_unpushed_commit_count = AsyncMock(return_value=3)
service._git.push_current_branch = AsyncMock(return_value=True)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
event = await gen.asend(None)
events.append(event)
assert isinstance(event, TeleportCheckingGitEvent)
event = await gen.asend(None)
events.append(event)
assert isinstance(event, TeleportPushRequiredEvent)
assert event.unpushed_count == 3
event = await gen.asend(TeleportPushResponseEvent(approved=True))
events.append(event)
assert isinstance(event, TeleportPushingEvent)
async for event in gen:
events.append(event)
assert isinstance(events[-1], TeleportCompleteEvent)
@pytest.mark.asyncio
async def test_execute_requires_push_and_user_declines(
self, service: TeleportService, git_info: GitRepoInfo
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=False)
service._git.get_unpushed_commit_count = AsyncMock(return_value=1)
session = TeleportSession()
gen = service.execute("test prompt", session)
await gen.asend(None)
await gen.asend(None)
with pytest.raises(ServiceTeleportError, match="Teleport cancelled"):
await gen.asend(TeleportPushResponseEvent(approved=False))
@pytest.mark.asyncio
async def test_execute_requires_oauth_flow(
self, service: TeleportService, git_info: GitRepoInfo
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=True)
github_pending = MagicMock()
github_pending.connected = False
github_pending.oauth_url = "https://github.com/login/oauth"
github_pending.error = "Please connect GitHub"
github_pending.status = GitHubStatus.WAITING_FOR_OAUTH
github_connected = MagicMock()
github_connected.connected = True
github_connected.oauth_url = None
github_connected.error = None
github_connected.status = GitHubStatus.CONNECTED
async def _oauth_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield github_pending
yield github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _oauth_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
events = []
gen = service.execute("test prompt", session)
async for event in gen:
events.append(event)
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
assert isinstance(events[2], TeleportWaitingForGitHubEvent)
assert events[2].message is None
assert isinstance(events[3], TeleportAuthRequiredEvent)
assert events[3].oauth_url == "https://github.com/login/oauth"
assert isinstance(events[4], TeleportWaitingForGitHubEvent)
assert events[4].message == "Please connect GitHub"
assert isinstance(events[5], TeleportAuthCompleteEvent)
assert isinstance(events[-1], TeleportCompleteEvent)
@pytest.mark.asyncio
async def test_execute_raises_when_chat_url_is_none(
self,
service: TeleportService,
git_info: GitRepoInfo,
mock_github_connected: MagicMock,
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=True)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(return_value=None)
service._nuage_client_instance = mock_nuage
session = TeleportSession()
gen = service.execute("test prompt", session)
with pytest.raises(ServiceTeleportError, match="not available"):
async for _ in gen:
pass
@pytest.mark.asyncio
async def test_execute_uses_default_prompt_when_none(
self,
service: TeleportService,
git_info: GitRepoInfo,
mock_github_connected: MagicMock,
) -> None:
service._git.get_info = AsyncMock(return_value=git_info)
service._git.is_commit_pushed = AsyncMock(return_value=True)
async def _connected_gen(*_a: object, **_kw: object): # type: ignore[no-untyped-def]
yield mock_github_connected
mock_nuage = MagicMock()
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
mock_nuage.wait_for_github_connection = _connected_gen
mock_nuage.get_chat_assistant_url = AsyncMock(
return_value="https://chat.example.com/123"
)
service._nuage_client_instance = mock_nuage
session = TeleportSession(
messages=[{"role": "user", "content": "help me refactor"}]
)
gen = service.execute(None, session)
async for _ in gen:
pass
call_args = mock_nuage.start_workflow.call_args
assert "teleported" in call_args[0][0].prompt.lower()
@pytest.mark.asyncio
async def test_execute_uses_experimental_nuage_when_enabled(
self, tmp_path: Path
) -> None:
seen_body: dict[str, object] | None = None
seen_url: str | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body, seen_url
seen_url = str(request.url)
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://chat.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
client=client,
)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="some local diff",
)
)
service._git.is_commit_pushed = AsyncMock(return_value=True)
service._git.is_branch_pushed = AsyncMock(return_value=True)
events = [
event
async for event in service.execute("test prompt", TeleportSession())
]
assert isinstance(events[0], TeleportCheckingGitEvent)
assert isinstance(events[1], TeleportStartingWorkflowEvent)
assert isinstance(events[2], TeleportCompleteEvent)
assert (
events[2].url == "https://chat.example.com/code/project-id/web-session-id"
)
assert service._nuage_client_instance is None
assert seen_url == "https://chat.example.com/api/v1/code/sessions"
assert seen_body is not None
assert seen_body["message"] == {
"role": "user",
"parts": [{"type": "text", "text": "test prompt"}],
}
repos = seen_body["context"]["repositories"]
assert len(repos) == 1
assert repos[0]["repoUrl"] == "https://github.com/owner/repo"
assert repos[0]["branch"] == "main"
assert repos[0]["commitSha"] == "abc123"
assert repos[0]["diff"]["format"] == "git-diff"
assert repos[0]["diff"]["encoding"] == "base64"
assert repos[0]["diff"]["compression"] == "zstd"
assert len(repos[0]["diff"]["content"]) > 0
assert "idempotencyKey" in seen_body
@pytest.mark.asyncio
async def test_execute_experimental_nuage_uses_last_message_when_prompt_missing(
self, tmp_path: Path
) -> None:
seen_body: dict[str, object] | None = None
async def handler(request: httpx.Request) -> httpx.Response:
nonlocal seen_body
seen_body = json.loads(request.content)
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
client=client,
)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
)
service._git.is_commit_pushed = AsyncMock(return_value=True)
service._git.is_branch_pushed = AsyncMock(return_value=True)
session = TeleportSession(
messages=[{"role": "user", "content": "help me refactor"}]
)
events = [event async for event in service.execute(None, session)]
assert isinstance(events[-1], TeleportCompleteEvent)
assert seen_body is not None
assert seen_body["message"] == {
"role": "user",
"parts": [{"type": "text", "text": "help me refactor (continue)"}],
}
@pytest.mark.asyncio
async def test_execute_experimental_nuage_requires_branch(
self, tmp_path: Path
) -> None:
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch=None,
commit="abc123",
diff="",
)
)
with pytest.raises(ServiceTeleportError, match="checked-out branch"):
async for _ in service.execute("test prompt", TeleportSession()):
pass
service._git.fetch.assert_not_awaited()
@pytest.mark.asyncio
async def test_execute_experimental_nuage_keeps_push_confirmation(
self, tmp_path: Path
) -> None:
async def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"sessionId": "controller-session-id",
"webSessionId": "web-session-id",
"projectId": "project-id",
"status": "running",
"url": "https://chat.example.com/code/project-id/web-session-id",
},
)
config = VibeConfig(vibe_code_experimental_nuage_enabled=True)
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
service = TeleportService(
session_logger=MagicMock(),
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
vibe_config=config,
client=client,
)
service._git.fetch = AsyncMock()
service._git.get_info = AsyncMock(
return_value=GitRepoInfo(
remote_url="https://github.com/owner/repo",
owner="owner",
repo="repo",
branch="main",
commit="abc123",
diff="",
)
)
service._git.is_commit_pushed = AsyncMock(return_value=False)
service._git.is_branch_pushed = AsyncMock(return_value=False)
service._git.get_unpushed_commit_count = AsyncMock(return_value=2)
service._git.push_current_branch = AsyncMock(return_value=True)
gen = service.execute("test prompt", TeleportSession())
assert isinstance(await gen.asend(None), TeleportCheckingGitEvent)
push_event = await gen.asend(None)
assert isinstance(push_event, TeleportPushRequiredEvent)
assert push_event.unpushed_count == 2
assert push_event.branch_not_pushed is True
assert isinstance(
await gen.asend(TeleportPushResponseEvent(approved=True)),
TeleportPushingEvent,
)
events = [event async for event in gen]
service._git.push_current_branch.assert_awaited_once()
assert isinstance(events[0], TeleportStartingWorkflowEvent)
assert isinstance(events[1], TeleportCompleteEvent)
class TestTeleportServiceContextManager:
@pytest.mark.asyncio
async def test_creates_client_on_enter(self, tmp_path: Path) -> None:
mock_session_logger = MagicMock()
service = TeleportService(
session_logger=mock_session_logger,
vibe_code_base_url="https://api.example.com",
vibe_code_workflow_id="workflow-id",
vibe_code_api_key="api-key",
workdir=tmp_path,
)
assert service._client is None
async with service:
assert service._client is not None
assert service._nuage_client_instance is not None
assert service._client is None
class TestTeleportAvailability:
def test_teleport_available_is_false_when_git_not_installed(self) -> None:
with patch.dict(os.environ, {"GIT_PYTHON_GIT_EXECUTABLE": "/nonexistent/git"}):
agent_loop = _reimport_agent_loop()
assert agent_loop._TELEPORT_AVAILABLE is False
def test_teleport_service_raises_error_when_git_not_available(self) -> None:
with patch.dict(os.environ, {"GIT_PYTHON_GIT_EXECUTABLE": "/nonexistent/git"}):
agent_loop = _reimport_agent_loop()
with pytest.raises(agent_loop.TeleportError, match="git to be installed"):
agent_loop.AgentLoop.teleport_service.fget(
MagicMock(_teleport_service=None)
)
def test_teleport_available_is_true_when_git_installed(
self, tmp_path: Path
) -> None:
fake_git = tmp_path / "git"
fake_git.write_text("#!/bin/sh\necho 'git version 2.0.0'")
fake_git.chmod(0o755)
with patch.dict(os.environ, {"GIT_PYTHON_GIT_EXECUTABLE": str(fake_git)}):
agent_loop = _reimport_agent_loop()
assert agent_loop._TELEPORT_AVAILABLE is True