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>
725 lines
27 KiB
Python
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
|