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
|
|
@ -2,9 +2,9 @@ from __future__ import annotations
|
|||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import build_test_vibe_config
|
||||
from vibe.core.agents.manager import AgentManager
|
||||
from vibe.core.agents.models import BUILTIN_AGENTS, EXPLORE, AgentSafety, AgentType
|
||||
from vibe.core.config import VibeConfig
|
||||
|
||||
|
||||
class TestAgentProfile:
|
||||
|
|
@ -31,7 +31,9 @@ class TestAgentProfile:
|
|||
class TestAgentManager:
|
||||
@pytest.fixture
|
||||
def manager(self) -> AgentManager:
|
||||
config = VibeConfig(include_project_context=False, include_prompt_detail=False)
|
||||
config = build_test_vibe_config(
|
||||
include_project_context=False, include_prompt_detail=False
|
||||
)
|
||||
return AgentManager(lambda: config)
|
||||
|
||||
def test_get_subagents_returns_only_subagents(self, manager: AgentManager) -> None:
|
||||
|
|
|
|||
46
tests/core/test_auth_crypto.py
Normal file
46
tests/core/test_auth_crypto.py
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict
|
||||
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
|
||||
from vibe.core.auth import EncryptedPayload, decrypt, encrypt
|
||||
|
||||
|
||||
def _generate_test_key_pair() -> tuple[bytes, bytes]:
|
||||
private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
|
||||
private_pem = private_key.private_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PrivateFormat.PKCS8,
|
||||
encryption_algorithm=serialization.NoEncryption(),
|
||||
)
|
||||
public_pem = private_key.public_key().public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
return private_pem, public_pem
|
||||
|
||||
|
||||
class TestEncryptDecrypt:
|
||||
def test_encrypt_decrypt_roundtrip(self) -> None:
|
||||
private_pem, public_pem = _generate_test_key_pair()
|
||||
|
||||
plaintext = "ghp_test_token_12345"
|
||||
encrypted = encrypt(plaintext, public_pem)
|
||||
|
||||
assert encrypted.encrypted_key != plaintext
|
||||
assert encrypted.ciphertext != plaintext
|
||||
|
||||
decrypted = decrypt(encrypted, private_pem)
|
||||
assert decrypted == plaintext
|
||||
|
||||
def test_encrypted_payload_serialization(self) -> None:
|
||||
payload = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce123", ciphertext="cipher"
|
||||
)
|
||||
|
||||
data = asdict(payload)
|
||||
restored = EncryptedPayload(**data)
|
||||
|
||||
assert restored == payload
|
||||
286
tests/core/test_auth_github.py
Normal file
286
tests/core/test_auth_github.py
Normal file
|
|
@ -0,0 +1,286 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from vibe.core.auth.github import (
|
||||
DeviceFlowHandle,
|
||||
DeviceFlowInfo,
|
||||
GitHubAuthError,
|
||||
GitHubAuthProvider,
|
||||
)
|
||||
|
||||
|
||||
class TestDeviceFlowModels:
|
||||
def test_device_flow_info(self) -> None:
|
||||
info = DeviceFlowInfo(
|
||||
user_code="ABC-123", verification_uri="https://example.com"
|
||||
)
|
||||
assert info.user_code == "ABC-123"
|
||||
assert info.verification_uri == "https://example.com"
|
||||
|
||||
def test_device_flow_handle(self) -> None:
|
||||
info = DeviceFlowInfo(
|
||||
user_code="ABC-123", verification_uri="https://example.com"
|
||||
)
|
||||
handle = DeviceFlowHandle(device_code="dc_123", expires_in=900, info=info)
|
||||
assert handle.device_code == "dc_123"
|
||||
assert handle.expires_in == 900
|
||||
assert handle.info.user_code == "ABC-123"
|
||||
|
||||
|
||||
class TestGitHubAuthProviderContextManager:
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_client_on_enter(self) -> None:
|
||||
provider = GitHubAuthProvider()
|
||||
assert provider._client is None
|
||||
async with provider:
|
||||
assert provider._client is not None
|
||||
assert provider._client is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_provided_client(self) -> None:
|
||||
external_client = httpx.AsyncClient()
|
||||
provider = GitHubAuthProvider(client=external_client)
|
||||
async with provider:
|
||||
assert provider._client is external_client
|
||||
assert provider._client is external_client
|
||||
await external_client.aclose()
|
||||
|
||||
|
||||
class TestGitHubAuthProviderGetToken:
|
||||
def test_returns_token_from_keyring(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
mock_keyring.get_password.return_value = "ghp_test_token"
|
||||
provider = GitHubAuthProvider()
|
||||
token = provider.get_token()
|
||||
assert token == "ghp_test_token"
|
||||
|
||||
def test_returns_none_on_keyring_error(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
import keyring.errors
|
||||
|
||||
mock_keyring.errors = keyring.errors
|
||||
mock_keyring.get_password.side_effect = keyring.errors.KeyringError("error")
|
||||
provider = GitHubAuthProvider()
|
||||
token = provider.get_token()
|
||||
assert token is None
|
||||
|
||||
def test_returns_none_when_no_token(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
mock_keyring.get_password.return_value = None
|
||||
provider = GitHubAuthProvider()
|
||||
token = provider.get_token()
|
||||
assert token is None
|
||||
|
||||
|
||||
class TestGitHubAuthProviderHasToken:
|
||||
def test_returns_true_when_token_exists(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
mock_keyring.get_password.return_value = "ghp_token"
|
||||
provider = GitHubAuthProvider()
|
||||
assert provider.has_token() is True
|
||||
|
||||
def test_returns_false_when_no_token(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
mock_keyring.get_password.return_value = None
|
||||
provider = GitHubAuthProvider()
|
||||
assert provider.has_token() is False
|
||||
|
||||
|
||||
class TestGitHubAuthProviderStartDeviceFlow:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
@pytest.fixture
|
||||
def provider(self, mock_client: MagicMock) -> GitHubAuthProvider:
|
||||
return GitHubAuthProvider(client=mock_client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_device_flow_success(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {
|
||||
"device_code": "dc_123",
|
||||
"user_code": "ABC-123",
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"expires_in": 900,
|
||||
}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.webbrowser") as mock_browser:
|
||||
handle = await provider.start_device_flow(open_browser=True)
|
||||
mock_browser.open.assert_called_once_with("https://github.com/login/device")
|
||||
|
||||
assert handle.device_code == "dc_123"
|
||||
assert handle.info.user_code == "ABC-123"
|
||||
assert handle.expires_in == 900
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_device_flow_without_browser(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {
|
||||
"device_code": "dc_123",
|
||||
"user_code": "ABC-123",
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"expires_in": 900,
|
||||
}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.webbrowser") as mock_browser:
|
||||
await provider.start_device_flow(open_browser=False)
|
||||
mock_browser.open.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_device_flow_failure(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Bad request"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with pytest.raises(GitHubAuthError, match="Failed to initiate device flow"):
|
||||
await provider.start_device_flow()
|
||||
|
||||
|
||||
class TestGitHubAuthProviderPollForToken:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
@pytest.fixture
|
||||
def provider(self, mock_client: MagicMock) -> GitHubAuthProvider:
|
||||
return GitHubAuthProvider(client=mock_client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_returns_token_on_success(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"access_token": "ghp_new_token"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock):
|
||||
token = await provider._poll_for_token(
|
||||
mock_client, "dc_123", expires_in=10, interval=1
|
||||
)
|
||||
assert token == "ghp_new_token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_handles_slow_down(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
responses = [
|
||||
MagicMock(
|
||||
json=MagicMock(return_value={"error": "slow_down", "interval": 5})
|
||||
),
|
||||
MagicMock(json=MagicMock(return_value={"access_token": "ghp_token"})),
|
||||
]
|
||||
mock_client.post = AsyncMock(side_effect=responses)
|
||||
|
||||
with patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock):
|
||||
token = await provider._poll_for_token(
|
||||
mock_client, "dc_123", expires_in=30, interval=1
|
||||
)
|
||||
assert token == "ghp_token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_raises_on_expired_token(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"error": "expired_token"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock):
|
||||
with pytest.raises(GitHubAuthError, match="expired_token"):
|
||||
await provider._poll_for_token(
|
||||
mock_client, "dc_123", expires_in=10, interval=1
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_raises_on_access_denied(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"error": "access_denied"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock):
|
||||
with pytest.raises(GitHubAuthError, match="access_denied"):
|
||||
await provider._poll_for_token(
|
||||
mock_client, "dc_123", expires_in=10, interval=1
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_raises_on_timeout(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"error": "authorization_pending"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock):
|
||||
with pytest.raises(GitHubAuthError, match="timed out"):
|
||||
await provider._poll_for_token(
|
||||
mock_client, "dc_123", expires_in=2, interval=1
|
||||
)
|
||||
|
||||
|
||||
class TestGitHubAuthProviderSaveToken:
|
||||
def test_save_token_success(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
provider = GitHubAuthProvider()
|
||||
provider._save_token("ghp_token")
|
||||
mock_keyring.set_password.assert_called_once_with(
|
||||
"vibe", "github_token", "ghp_token"
|
||||
)
|
||||
|
||||
def test_save_token_raises_on_keyring_error(self) -> None:
|
||||
with patch("vibe.core.auth.github.keyring") as mock_keyring:
|
||||
import keyring.errors
|
||||
|
||||
mock_keyring.errors = keyring.errors
|
||||
mock_keyring.set_password.side_effect = keyring.errors.KeyringError("error")
|
||||
provider = GitHubAuthProvider()
|
||||
with pytest.raises(GitHubAuthError, match="Failed to save token"):
|
||||
provider._save_token("ghp_token")
|
||||
|
||||
|
||||
class TestGitHubAuthProviderWaitForToken:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
@pytest.fixture
|
||||
def provider(self, mock_client: MagicMock) -> GitHubAuthProvider:
|
||||
return GitHubAuthProvider(client=mock_client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wait_for_token_polls_and_saves(
|
||||
self, provider: GitHubAuthProvider, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"access_token": "ghp_token"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
info = DeviceFlowInfo(user_code="ABC", verification_uri="https://example.com")
|
||||
handle = DeviceFlowHandle(device_code="dc_123", expires_in=10, info=info)
|
||||
|
||||
with (
|
||||
patch("vibe.core.auth.github.asyncio.sleep", new_callable=AsyncMock),
|
||||
patch("vibe.core.auth.github.keyring") as mock_keyring,
|
||||
):
|
||||
token = await provider.wait_for_token(handle)
|
||||
|
||||
assert token == "ghp_token"
|
||||
mock_keyring.set_password.assert_called_once()
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
import threading
|
||||
|
||||
from vibe.core.config import load_dotenv_values
|
||||
|
||||
|
|
@ -59,3 +61,21 @@ def test_ignores_empty_values(tmp_path: Path) -> None:
|
|||
assert "EMPTY" not in environ
|
||||
assert "MISTRAL_API_KEY" not in environ
|
||||
assert "NO_VALUE" not in environ
|
||||
|
||||
|
||||
def test_reads_from_fifo(tmp_path: Path) -> None:
|
||||
fifo_path = tmp_path / ".env.fifo"
|
||||
os.mkfifo(fifo_path)
|
||||
environ: dict[str, str] = {}
|
||||
|
||||
def write_to_fifo() -> None:
|
||||
with open(fifo_path, "w") as f:
|
||||
f.write("FIFO_KEY=fifo-value\n")
|
||||
|
||||
writer = threading.Thread(target=write_to_fifo)
|
||||
writer.start()
|
||||
|
||||
load_dotenv_values(env_path=fifo_path, environ=environ)
|
||||
|
||||
writer.join()
|
||||
assert environ["FIFO_KEY"] == "fifo-value"
|
||||
|
|
|
|||
327
tests/core/test_teleport_git.py
Normal file
327
tests/core/test_teleport_git.py
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.teleport.errors import (
|
||||
ServiceTeleportError,
|
||||
ServiceTeleportNotSupportedError,
|
||||
)
|
||||
from vibe.core.teleport.git import GitRepoInfo, GitRepository
|
||||
|
||||
|
||||
def make_mock_remote(url: str) -> MagicMock:
|
||||
remote = MagicMock()
|
||||
remote.urls = [url]
|
||||
return remote
|
||||
|
||||
|
||||
def make_mock_repo(
|
||||
urls: list[str] | None = None,
|
||||
commit: str | None = "abc123",
|
||||
branch: str | None = "main",
|
||||
is_detached: bool = False,
|
||||
diff: str = "",
|
||||
) -> MagicMock:
|
||||
mock = MagicMock()
|
||||
if urls:
|
||||
mock.remotes = [make_mock_remote(url) for url in urls]
|
||||
else:
|
||||
mock.remotes = []
|
||||
mock.head.commit.hexsha = commit
|
||||
mock.head.is_detached = is_detached
|
||||
mock.active_branch.name = branch
|
||||
mock.git.diff.return_value = diff
|
||||
mock.git.branch.return_value = ""
|
||||
mock.git.rev_list.return_value = "0"
|
||||
mock.git.rev_parse.return_value = "abc123"
|
||||
mock.remote.return_value = make_mock_remote(urls[0] if urls else "")
|
||||
return mock
|
||||
|
||||
|
||||
class TestGitRepositoryParseGithubUrl:
|
||||
def test_parse_ssh_url(self) -> None:
|
||||
result = GitRepository._parse_github_url("git@github.com:owner/repo.git")
|
||||
assert result == ("owner", "repo")
|
||||
|
||||
def test_parse_ssh_url_without_git_suffix(self) -> None:
|
||||
result = GitRepository._parse_github_url("git@github.com:owner/repo")
|
||||
assert result == ("owner", "repo")
|
||||
|
||||
def test_parse_https_url(self) -> None:
|
||||
result = GitRepository._parse_github_url("https://github.com/owner/repo.git")
|
||||
assert result == ("owner", "repo")
|
||||
|
||||
def test_parse_https_url_without_git_suffix(self) -> None:
|
||||
result = GitRepository._parse_github_url("https://github.com/owner/repo")
|
||||
assert result == ("owner", "repo")
|
||||
|
||||
def test_parse_https_url_with_credentials(self) -> None:
|
||||
result = GitRepository._parse_github_url(
|
||||
"https://x-access-token:gho_xxxx@github.com/owner/repo.git"
|
||||
)
|
||||
assert result == ("owner", "repo")
|
||||
|
||||
def test_parse_non_github_url_returns_none(self) -> None:
|
||||
result = GitRepository._parse_github_url("git@gitlab.com:owner/repo.git")
|
||||
assert result is None
|
||||
|
||||
def test_parse_invalid_url_returns_none(self) -> None:
|
||||
result = GitRepository._parse_github_url("not-a-valid-url")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestGitRepositoryToHttpsUrl:
|
||||
def test_converts_to_https_url(self) -> None:
|
||||
result = GitRepository._to_https_url("owner", "repo")
|
||||
assert result == "https://github.com/owner/repo.git"
|
||||
|
||||
|
||||
class TestGitRepositoryIsSupported:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_not_git_repo(self, tmp_path: Path) -> None:
|
||||
repo = GitRepository(tmp_path)
|
||||
assert await repo.is_supported() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_no_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=None)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_supported() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_non_github_remote(
|
||||
self, repo: GitRepository
|
||||
) -> None:
|
||||
mock = make_mock_repo(urls=["git@gitlab.com:owner/repo.git"])
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_supported() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_true_when_github_repo(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"])
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_supported() is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finds_github_among_multiple_remotes(
|
||||
self, repo: GitRepository
|
||||
) -> None:
|
||||
mock = make_mock_repo(
|
||||
urls=["git@gitlab.com:owner/repo.git", "git@github.com:owner/repo.git"]
|
||||
)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_supported() is True
|
||||
|
||||
|
||||
class TestGitRepositoryGetInfo:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_not_git_repo(self, tmp_path: Path) -> None:
|
||||
repo = GitRepository(tmp_path)
|
||||
with pytest.raises(
|
||||
ServiceTeleportNotSupportedError, match="Not a git repository"
|
||||
):
|
||||
await repo.get_info()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_no_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=None)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
with pytest.raises(
|
||||
ServiceTeleportNotSupportedError, match="No GitHub remote"
|
||||
):
|
||||
await repo.get_info()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_non_github_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@gitlab.com:owner/repo.git"])
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
with pytest.raises(
|
||||
ServiceTeleportNotSupportedError, match="No GitHub remote"
|
||||
):
|
||||
await repo.get_info()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_no_commit(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], commit=None)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
with pytest.raises(
|
||||
ServiceTeleportNotSupportedError,
|
||||
match="Could not determine current commit",
|
||||
):
|
||||
await repo.get_info()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_info_on_success(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(
|
||||
urls=["git@github.com:owner/repo.git"],
|
||||
commit="abc123def456",
|
||||
branch="main",
|
||||
diff="diff content",
|
||||
)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
info = await repo.get_info()
|
||||
assert info == GitRepoInfo(
|
||||
remote_url="https://github.com/owner/repo.git",
|
||||
owner="owner",
|
||||
repo="repo",
|
||||
branch="main",
|
||||
commit="abc123def456",
|
||||
diff="diff content",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_detached_head(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(
|
||||
urls=["git@github.com:owner/repo.git"],
|
||||
commit="abc123def456",
|
||||
is_detached=True,
|
||||
)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
info = await repo.get_info()
|
||||
assert info.branch is None
|
||||
|
||||
|
||||
class TestGitRepositoryIsCommitPushed:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_not_on_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"])
|
||||
mock.git.branch.return_value = ""
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_commit_pushed("abc123") is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_true_when_on_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"])
|
||||
mock.git.branch.return_value = " origin/main\n origin/feature"
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_commit_pushed("abc123") is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checks_correct_remote(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"])
|
||||
mock.git.branch.return_value = " upstream/main\n"
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.is_commit_pushed("abc123", remote="upstream") is True
|
||||
assert await repo.is_commit_pushed("abc123", remote="origin") is False
|
||||
|
||||
|
||||
class TestGitRepositoryGetUnpushedCommitCount:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_no_branch(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], is_detached=True)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
with pytest.raises(ServiceTeleportError, match="no current branch"):
|
||||
await repo.get_unpushed_commit_count()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_count(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], branch="main")
|
||||
mock.git.rev_list.return_value = "5"
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.get_unpushed_commit_count() == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_to_default_branch(self, repo: GitRepository) -> None:
|
||||
from git.exc import GitCommandError
|
||||
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], branch="feature")
|
||||
|
||||
def rev_list_side_effect(*args: object, **kwargs: object) -> str:
|
||||
if "origin/feature..HEAD" in args:
|
||||
raise GitCommandError("git rev-list", 128)
|
||||
if "origin/main..HEAD" in args:
|
||||
return "3"
|
||||
return "0"
|
||||
|
||||
mock.git.rev_list.side_effect = rev_list_side_effect
|
||||
with (
|
||||
patch.object(repo, "_repo_or_raise", return_value=mock),
|
||||
patch.object(
|
||||
repo,
|
||||
"_get_remote_default_branch",
|
||||
AsyncMock(return_value="origin/main"),
|
||||
),
|
||||
):
|
||||
assert await repo.get_unpushed_commit_count() == 3
|
||||
|
||||
|
||||
class TestGitRepositoryPushCurrentBranch:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_no_branch(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], is_detached=True)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.push_current_branch() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_push_fails(self, repo: GitRepository) -> None:
|
||||
from git.exc import GitCommandError
|
||||
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], branch="main")
|
||||
mock.remote.return_value.push.side_effect = GitCommandError("git push", 1)
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.push_current_branch() is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_true_when_push_succeeds(self, repo: GitRepository) -> None:
|
||||
mock = make_mock_repo(urls=["git@github.com:owner/repo.git"], branch="main")
|
||||
with patch.object(repo, "_repo_or_raise", return_value=mock):
|
||||
assert await repo.push_current_branch() is True
|
||||
|
||||
|
||||
class TestGitRepositoryGetRemoteDefaultBranch:
|
||||
@pytest.fixture
|
||||
def repo(self, tmp_path: Path) -> GitRepository:
|
||||
return GitRepository(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_default_branch_when_head_exists(
|
||||
self, repo: GitRepository
|
||||
) -> None:
|
||||
mock = MagicMock()
|
||||
mock_head_ref = MagicMock()
|
||||
mock_head_ref.reference.name = "origin/main"
|
||||
mock.remotes.__getitem__.return_value.refs.HEAD = mock_head_ref
|
||||
mock.git.rev_parse.return_value = "abc123"
|
||||
result = await repo._get_remote_default_branch(mock, "origin")
|
||||
assert result == "origin/main"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_none_when_remote_not_found(
|
||||
self, repo: GitRepository
|
||||
) -> None:
|
||||
mock = MagicMock()
|
||||
mock.remotes.__getitem__.side_effect = IndexError("No item found")
|
||||
result = await repo._get_remote_default_branch(mock, "upstream")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_none_when_head_ref_missing(
|
||||
self, repo: GitRepository
|
||||
) -> None:
|
||||
mock = MagicMock()
|
||||
mock.remotes.__getitem__.return_value.refs.HEAD = None
|
||||
result = await repo._get_remote_default_branch(mock, "origin")
|
||||
assert result is None
|
||||
337
tests/core/test_teleport_nuage.py
Normal file
337
tests/core/test_teleport_nuage.py
Normal file
|
|
@ -0,0 +1,337 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from vibe.core.auth import EncryptedPayload
|
||||
from vibe.core.teleport.errors import ServiceTeleportError
|
||||
from vibe.core.teleport.nuage import (
|
||||
CreateLeChatThreadInput,
|
||||
GitRepoConfig,
|
||||
NuageClient,
|
||||
VibeNewSandbox,
|
||||
VibeSandboxConfig,
|
||||
WorkflowParams,
|
||||
)
|
||||
|
||||
|
||||
class TestNuageModels:
|
||||
def test_git_repo_config_defaults(self) -> None:
|
||||
config = GitRepoConfig(url="https://github.com/owner/repo.git")
|
||||
assert config.url == "https://github.com/owner/repo.git"
|
||||
assert config.branch is None
|
||||
assert config.commit is None
|
||||
|
||||
def test_git_repo_config_with_values(self) -> None:
|
||||
config = GitRepoConfig(
|
||||
url="https://github.com/owner/repo.git", branch="main", commit="abc123"
|
||||
)
|
||||
assert config.branch == "main"
|
||||
assert config.commit == "abc123"
|
||||
|
||||
def test_vibe_sandbox_config_defaults(self) -> None:
|
||||
config = VibeSandboxConfig()
|
||||
assert config.git_repo is None
|
||||
|
||||
def test_vibe_new_sandbox_defaults(self) -> None:
|
||||
sandbox = VibeNewSandbox()
|
||||
assert sandbox.type == "new"
|
||||
assert sandbox.config.git_repo is None
|
||||
assert sandbox.teleported_diffs is None
|
||||
|
||||
def test_workflow_params_serialization(self) -> None:
|
||||
params = WorkflowParams(
|
||||
prompt="test prompt",
|
||||
sandbox=VibeNewSandbox(
|
||||
config=VibeSandboxConfig(
|
||||
git_repo=GitRepoConfig(
|
||||
url="https://github.com/owner/repo.git",
|
||||
branch="main",
|
||||
commit="abc123",
|
||||
)
|
||||
),
|
||||
teleported_diffs=b"base64data",
|
||||
),
|
||||
)
|
||||
data = params.model_dump()
|
||||
assert data["prompt"] == "test prompt"
|
||||
assert data["sandbox"]["type"] == "new"
|
||||
assert data["sandbox"]["config"]["git_repo"]["url"] == (
|
||||
"https://github.com/owner/repo.git"
|
||||
)
|
||||
assert data["sandbox"]["teleported_diffs"] == b"base64data"
|
||||
|
||||
def test_create_le_chat_thread_input(self) -> None:
|
||||
input_data = CreateLeChatThreadInput(
|
||||
encrypted_api_key={"key": "value"},
|
||||
user_message="test message",
|
||||
project_name="test-project",
|
||||
)
|
||||
assert input_data.encrypted_api_key == {"key": "value"}
|
||||
assert input_data.user_message == "test message"
|
||||
assert input_data.project_name == "test-project"
|
||||
|
||||
|
||||
class TestNuageClientContextManager:
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_client_on_enter(self) -> None:
|
||||
nuage = NuageClient("https://api.example.com", "api-key", "workflow-id")
|
||||
assert nuage._client is None
|
||||
async with nuage:
|
||||
assert nuage._client is not None
|
||||
assert nuage._client is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_provided_client(self) -> None:
|
||||
external_client = httpx.AsyncClient()
|
||||
nuage = NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=external_client
|
||||
)
|
||||
async with nuage:
|
||||
assert nuage._client is external_client
|
||||
assert nuage._client is external_client
|
||||
await external_client.aclose()
|
||||
|
||||
|
||||
class TestNuageClientStartWorkflow:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
client = MagicMock(spec=httpx.AsyncClient)
|
||||
return client
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_workflow_success(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = True
|
||||
mock_response.json.return_value = {"execution_id": "exec-123"}
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
execution_id = await nuage.start_workflow(params)
|
||||
|
||||
assert execution_id == "exec-123"
|
||||
mock_client.post.assert_called_once()
|
||||
call_args = mock_client.post.call_args
|
||||
assert "workflows/workflow-id/execute" in call_args[0][0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_workflow_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Internal Server Error"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
with pytest.raises(ServiceTeleportError, match="Nuage workflow trigger failed"):
|
||||
await nuage.start_workflow(params)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_workflow_unauthorized_hint(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Unauthorized"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
params = WorkflowParams(prompt="test", sandbox=VibeNewSandbox())
|
||||
with pytest.raises(ServiceTeleportError, match="STAGING_MISTRAL_API_KEY"):
|
||||
await nuage.start_workflow(params)
|
||||
|
||||
|
||||
class TestNuageClientSendGithubToken:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_github_token_success(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
signal_response = MagicMock()
|
||||
signal_response.is_success = True
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, signal_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
|
||||
await nuage.send_github_token("exec-123", "ghp_token")
|
||||
|
||||
assert mock_client.post.call_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_public_key_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Not found"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
with pytest.raises(ServiceTeleportError, match="Failed to get public key"):
|
||||
await nuage._query_public_key("exec-123")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_signal_encrypted_token_failure(
|
||||
self, nuage: NuageClient, mock_client: MagicMock
|
||||
) -> None:
|
||||
mock_response = MagicMock()
|
||||
mock_response.is_success = False
|
||||
mock_response.text = "Signal failed"
|
||||
mock_client.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
with pytest.raises(ServiceTeleportError, match="Failed to send GitHub token"):
|
||||
await nuage._signal_encrypted_token("exec-123", encrypted)
|
||||
|
||||
|
||||
class TestNuageClientCreateLeChatThread:
|
||||
@pytest.fixture
|
||||
def mock_client(self) -> MagicMock:
|
||||
return MagicMock(spec=httpx.AsyncClient)
|
||||
|
||||
@pytest.fixture
|
||||
def nuage(self, mock_client: MagicMock) -> NuageClient:
|
||||
return NuageClient(
|
||||
"https://api.example.com", "api-key", "workflow-id", client=mock_client
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_le_chat_thread_success(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
update_response = MagicMock()
|
||||
update_response.is_success = True
|
||||
update_response.json.return_value = {
|
||||
"result": {"chat_url": "https://chat.example.com/thread/123"}
|
||||
}
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, update_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
|
||||
url = await nuage.create_le_chat_thread("exec-123", "test message")
|
||||
assert url == "https://chat.example.com/thread/123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_le_chat_thread_failure(
|
||||
self,
|
||||
nuage: NuageClient,
|
||||
mock_client: MagicMock,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
query_response = MagicMock()
|
||||
query_response.is_success = True
|
||||
query_response.json.return_value = {
|
||||
"result": {
|
||||
"public_key": (
|
||||
"-----BEGIN PUBLIC KEY-----\n"
|
||||
"MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA0Z3VS5JJcds3xfn/"
|
||||
"ygWyf8TFXQNZ0XsLOqXB1Mi2+bKPFv1WfhECTxJ3c6SXl0p1sGyWTFxRV8u/"
|
||||
"bKqYZ0E6VZ6YRTRPFiGq0kkONjVBFxOQ8Y0jeT0d9e0Y3E3MWDL8tQ0Nz9v8"
|
||||
"5Y7gC8F1m/dEbBwPjCJQV0Dg0z3gZDO8RCG0GrBoLO0b+NNqL8FXPPDXQ1l4"
|
||||
"FGnYM0gZ1rCU7Y/zTN1wI4sCQ0GJQPDA1hWB8KRJl5x0ZDXE3rRwT1E8c+Fn"
|
||||
"ZFV1nN0C6zxF7GpVY3FVWXS4PA0FH+8C1+TnYgBL7xS0o+LF6PgjGT5F3CXD"
|
||||
"BZmYSxKL+EsVVGT5EuYbJE9TxVwIDAQAB\n"
|
||||
"-----END PUBLIC KEY-----"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
update_response = MagicMock()
|
||||
update_response.is_success = False
|
||||
update_response.text = "Failed to create thread"
|
||||
|
||||
mock_client.post = AsyncMock(side_effect=[query_response, update_response])
|
||||
|
||||
encrypted = EncryptedPayload(
|
||||
encrypted_key="enc_key", nonce="nonce", ciphertext="cipher"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vibe.core.teleport.nuage.encrypt", lambda _token, _key: encrypted
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ServiceTeleportError, match="Failed to create Le Chat thread"
|
||||
):
|
||||
await nuage.create_le_chat_thread("exec-123", "test message")
|
||||
|
||||
|
||||
class TestNuageClientHeaders:
|
||||
def test_headers_include_auth(self) -> None:
|
||||
nuage = NuageClient("https://api.example.com", "test-api-key", "workflow-id")
|
||||
headers = nuage._headers()
|
||||
assert headers["Authorization"] == "Bearer test-api-key"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
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