2.1.0 (#317)
Co-authored-by: Quentin Torroba <quentin.torroba@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: Clément Siriex <clement.sirieix@mistral.ai> Co-authored-by: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai> Co-authored-by: Nicolas Karolak <nicolas@karolak.fr>
This commit is contained in:
Mathias Gesbert
GitHub
parent
9809cfc831
commit
51fecc67d9
176 changed files with 8652 additions and 4451 deletions
418
tests/core/test_teleport_service.py
Normal file
418
tests/core/test_teleport_service.py
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import zstandard
|
||||
|
||||
from vibe.core.teleport.errors import (
|
||||
ServiceTeleportError,
|
||||
ServiceTeleportNotSupportedError,
|
||||
)
|
||||
from vibe.core.teleport.git import GitRepoInfo
|
||||
from vibe.core.teleport.nuage import TeleportSession
|
||||
from vibe.core.teleport.teleport import _NUAGE_EXECUTION_URL_TEMPLATE, TeleportService
|
||||
from vibe.core.teleport.types import (
|
||||
TeleportAuthCompleteEvent,
|
||||
TeleportAuthRequiredEvent,
|
||||
TeleportCheckingGitEvent,
|
||||
TeleportCompleteEvent,
|
||||
TeleportPushingEvent,
|
||||
TeleportPushRequiredEvent,
|
||||
TeleportPushResponseEvent,
|
||||
TeleportSendingGithubTokenEvent,
|
||||
TeleportStartingWorkflowEvent,
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_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 TestTeleportServiceBuildSandbox:
|
||||
@pytest.fixture
|
||||
def service(self, tmp_path: Path) -> TeleportService:
|
||||
mock_session_logger = MagicMock()
|
||||
return TeleportService(
|
||||
session_logger=mock_session_logger,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="api-key",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
|
||||
def test_builds_sandbox_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="",
|
||||
)
|
||||
sandbox = service._build_sandbox(git_info)
|
||||
|
||||
assert sandbox.type == "new"
|
||||
assert sandbox.config.git_repo is not None
|
||||
assert sandbox.config.git_repo.url == "https://github.com/owner/repo.git"
|
||||
assert sandbox.config.git_repo.branch == "main"
|
||||
assert sandbox.config.git_repo.commit == "abc123"
|
||||
assert sandbox.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",
|
||||
)
|
||||
sandbox = service._build_sandbox(git_info)
|
||||
|
||||
assert sandbox.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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
with pytest.raises(
|
||||
ServiceTeleportError, match="STAGING_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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="valid-key",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_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("Not a git repository")
|
||||
)
|
||||
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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="api-key",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
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.mark.asyncio
|
||||
async def test_execute_happy_path_commit_pushed_with_token(
|
||||
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)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_existing_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
service._nuage = 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], TeleportSendingGithubTokenEvent)
|
||||
assert isinstance(events[3], TeleportCompleteEvent)
|
||||
expected_url = _NUAGE_EXECUTION_URL_TEMPLATE.format(
|
||||
workflow_id="workflow-id", execution_id="exec-123"
|
||||
)
|
||||
assert events[3].url == expected_url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_requires_push_and_user_approves(
|
||||
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=3)
|
||||
service._git.push_current_branch = AsyncMock(return_value=True)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
service._nuage = 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_auth_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)
|
||||
|
||||
mock_handle = MagicMock()
|
||||
mock_handle.info.user_code = "ABC-123"
|
||||
mock_handle.info.verification_uri = "https://github.com/login/device"
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value=None)
|
||||
mock_github_auth.start_device_flow = AsyncMock(return_value=mock_handle)
|
||||
mock_github_auth.wait_for_token = AsyncMock(return_value="ghp_new_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
service._nuage = 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], TeleportAuthRequiredEvent)
|
||||
assert events[1].user_code == "ABC-123"
|
||||
assert isinstance(events[2], TeleportAuthCompleteEvent)
|
||||
assert isinstance(events[3], TeleportStartingWorkflowEvent)
|
||||
assert isinstance(events[-1], TeleportCompleteEvent)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uses_default_prompt_when_none(
|
||||
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)
|
||||
|
||||
mock_github_auth = MagicMock()
|
||||
mock_github_auth.get_valid_token = AsyncMock(return_value="ghp_token")
|
||||
service._github_auth = mock_github_auth
|
||||
|
||||
mock_nuage = MagicMock()
|
||||
mock_nuage.start_workflow = AsyncMock(return_value="exec-123")
|
||||
mock_nuage.send_github_token = AsyncMock()
|
||||
service._nuage = mock_nuage
|
||||
|
||||
session = TeleportSession()
|
||||
gen = service.execute(None, session)
|
||||
async for _ in gen:
|
||||
pass
|
||||
|
||||
call_args = mock_nuage.start_workflow.call_args
|
||||
assert "continue where you left off" in call_args[0][0].prompt
|
||||
|
||||
|
||||
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,
|
||||
nuage_base_url="https://api.example.com",
|
||||
nuage_workflow_id="workflow-id",
|
||||
nuage_api_key="api-key",
|
||||
workdir=tmp_path,
|
||||
)
|
||||
assert service._client is None
|
||||
async with service:
|
||||
assert service._client is not None
|
||||
assert service._github_auth is not None
|
||||
assert service._nuage 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue