v2.16.1 (#808)
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai> Co-authored-by: Paul VEZIA <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin
GitHub
parent
c2cb612ac1
commit
564a14365e
22 changed files with 952 additions and 252 deletions
|
|
@ -1,18 +1,82 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.nuage.workflow import (
|
||||
WorkflowExecutionListResponse,
|
||||
WorkflowExecutionStatus,
|
||||
WorkflowExecutionWithoutResultResponse,
|
||||
)
|
||||
from vibe.core.session.resume_sessions import (
|
||||
RemoteResumeResult,
|
||||
RemoteResumeSessions,
|
||||
ResumeSessionInfo,
|
||||
can_delete_resume_session_source,
|
||||
list_remote_resume_sessions,
|
||||
session_latest_messages,
|
||||
short_session_id,
|
||||
)
|
||||
from vibe.core.session.session_id import shorten_session_id
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RemoteResumeRequest:
|
||||
workflow_identifier: str | None
|
||||
page_size: int
|
||||
status: Sequence[WorkflowExecutionStatus] | None
|
||||
|
||||
|
||||
class FakeRemoteResumeClient:
|
||||
def __init__(
|
||||
self,
|
||||
response: WorkflowExecutionListResponse | None = None,
|
||||
*,
|
||||
delay: float = 0.0,
|
||||
) -> None:
|
||||
self.response = response or WorkflowExecutionListResponse(executions=[])
|
||||
self.delay = delay
|
||||
self.requests: list[RemoteResumeRequest] = []
|
||||
self.closed = False
|
||||
|
||||
async def get_workflow_runs(
|
||||
self,
|
||||
workflow_identifier: str | None = None,
|
||||
page_size: int = 50,
|
||||
next_page_token: str | None = None,
|
||||
status: Sequence[WorkflowExecutionStatus] | None = None,
|
||||
user_id: str = "current",
|
||||
) -> WorkflowExecutionListResponse:
|
||||
self.requests.append(
|
||||
RemoteResumeRequest(
|
||||
workflow_identifier=workflow_identifier,
|
||||
page_size=page_size,
|
||||
status=status,
|
||||
)
|
||||
)
|
||||
if self.delay:
|
||||
await asyncio.sleep(self.delay)
|
||||
return self.response
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def enabled_vibe_code_config() -> MagicMock:
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = True
|
||||
config.vibe_code_api_key = "test-key"
|
||||
config.vibe_code_base_url = "https://test.example.com"
|
||||
config.api_timeout = 30
|
||||
config.vibe_code_workflow_id = "workflow-1"
|
||||
return config
|
||||
|
||||
|
||||
class TestShortenSessionId:
|
||||
def test_shortens_to_first_8_chars(self) -> None:
|
||||
sid = "abcdef1234567890"
|
||||
|
|
@ -67,40 +131,8 @@ class TestCanDeleteResumeSession:
|
|||
|
||||
|
||||
class TestListRemoteResumeSessions:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_when_vibe_code_disabled(self) -> None:
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = False
|
||||
config.vibe_code_api_key = "key"
|
||||
result = await list_remote_resume_sessions(config)
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_when_no_api_key(self) -> None:
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = True
|
||||
config.vibe_code_api_key = None
|
||||
result = await list_remote_resume_sessions(config)
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_when_both_missing(self) -> None:
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = False
|
||||
config.vibe_code_api_key = None
|
||||
result = await list_remote_resume_sessions(config)
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passes_active_statuses_to_api(self) -> None:
|
||||
from datetime import datetime
|
||||
|
||||
from vibe.core.nuage.workflow import (
|
||||
WorkflowExecutionListResponse,
|
||||
WorkflowExecutionStatus,
|
||||
WorkflowExecutionWithoutResultResponse,
|
||||
)
|
||||
|
||||
running = WorkflowExecutionWithoutResultResponse(
|
||||
workflow_name="vibe",
|
||||
execution_id="exec-running",
|
||||
|
|
@ -117,47 +149,28 @@ class TestListRemoteResumeSessions:
|
|||
)
|
||||
|
||||
mock_response = WorkflowExecutionListResponse(executions=[running, continued])
|
||||
client = FakeRemoteResumeClient(mock_response)
|
||||
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = True
|
||||
config.vibe_code_api_key = "test-key"
|
||||
config.vibe_code_base_url = "https://test.example.com"
|
||||
config.api_timeout = 30
|
||||
config.vibe_code_workflow_id = "workflow-1"
|
||||
|
||||
with patch("vibe.core.session.resume_sessions.WorkflowsClient") as MockClient:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get_workflow_runs.return_value = mock_response
|
||||
MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await list_remote_resume_sessions(config)
|
||||
result = await list_remote_resume_sessions(client, "workflow-1")
|
||||
|
||||
assert len(result) == 2
|
||||
session_ids = {s.session_id for s in result}
|
||||
assert "exec-running" in session_ids
|
||||
assert "exec-continued" in session_ids
|
||||
assert all(s.source == "remote" for s in result)
|
||||
|
||||
mock_client.get_workflow_runs.assert_called_once_with(
|
||||
workflow_identifier="workflow-1",
|
||||
page_size=50,
|
||||
status=[
|
||||
WorkflowExecutionStatus.RUNNING,
|
||||
WorkflowExecutionStatus.CONTINUED_AS_NEW,
|
||||
],
|
||||
)
|
||||
assert client.requests == [
|
||||
RemoteResumeRequest(
|
||||
workflow_identifier="workflow-1",
|
||||
page_size=50,
|
||||
status=[
|
||||
WorkflowExecutionStatus.RUNNING,
|
||||
WorkflowExecutionStatus.CONTINUED_AS_NEW,
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deduplicates_execution_ids_keeps_latest(self) -> None:
|
||||
from datetime import datetime
|
||||
|
||||
from vibe.core.nuage.workflow import (
|
||||
WorkflowExecutionListResponse,
|
||||
WorkflowExecutionStatus,
|
||||
WorkflowExecutionWithoutResultResponse,
|
||||
)
|
||||
|
||||
older = WorkflowExecutionWithoutResultResponse(
|
||||
workflow_name="vibe",
|
||||
execution_id="exec-1",
|
||||
|
|
@ -181,21 +194,9 @@ class TestListRemoteResumeSessions:
|
|||
)
|
||||
|
||||
mock_response = WorkflowExecutionListResponse(executions=[older, newer, other])
|
||||
client = FakeRemoteResumeClient(mock_response)
|
||||
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = True
|
||||
config.vibe_code_api_key = "test-key"
|
||||
config.vibe_code_base_url = "https://test.example.com"
|
||||
config.api_timeout = 30
|
||||
config.vibe_code_workflow_id = "workflow-1"
|
||||
|
||||
with patch("vibe.core.session.resume_sessions.WorkflowsClient") as MockClient:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get_workflow_runs.return_value = mock_response
|
||||
MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await list_remote_resume_sessions(config)
|
||||
result = await list_remote_resume_sessions(client, "workflow-1")
|
||||
|
||||
assert len(result) == 2
|
||||
by_id = {s.session_id: s for s in result}
|
||||
|
|
@ -206,14 +207,6 @@ class TestListRemoteResumeSessions:
|
|||
async def test_dedup_keeps_latest_start_time_when_previous_has_end_time(
|
||||
self,
|
||||
) -> None:
|
||||
from datetime import datetime
|
||||
|
||||
from vibe.core.nuage.workflow import (
|
||||
WorkflowExecutionListResponse,
|
||||
WorkflowExecutionStatus,
|
||||
WorkflowExecutionWithoutResultResponse,
|
||||
)
|
||||
|
||||
previous = WorkflowExecutionWithoutResultResponse(
|
||||
workflow_name="vibe",
|
||||
execution_id="exec-1",
|
||||
|
|
@ -230,22 +223,85 @@ class TestListRemoteResumeSessions:
|
|||
)
|
||||
|
||||
mock_response = WorkflowExecutionListResponse(executions=[previous, newer])
|
||||
client = FakeRemoteResumeClient(mock_response)
|
||||
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = True
|
||||
config.vibe_code_api_key = "test-key"
|
||||
config.vibe_code_base_url = "https://test.example.com"
|
||||
config.api_timeout = 30
|
||||
config.vibe_code_workflow_id = "workflow-1"
|
||||
|
||||
with patch("vibe.core.session.resume_sessions.WorkflowsClient") as MockClient:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.get_workflow_runs.return_value = mock_response
|
||||
MockClient.return_value.__aenter__ = AsyncMock(return_value=mock_client)
|
||||
MockClient.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await list_remote_resume_sessions(config)
|
||||
result = await list_remote_resume_sessions(client, "workflow-1")
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].session_id == "exec-1"
|
||||
assert result[0].status == WorkflowExecutionStatus.RUNNING
|
||||
|
||||
|
||||
class TestSessionLatestMessages:
|
||||
def test_remote_session_formats_title_and_status(self) -> None:
|
||||
session = ResumeSessionInfo(
|
||||
session_id="exec-1",
|
||||
source="remote",
|
||||
cwd="",
|
||||
title="My run",
|
||||
end_time=None,
|
||||
status="RUNNING",
|
||||
)
|
||||
messages = session_latest_messages([session], MagicMock())
|
||||
assert messages[session.option_id] == "My run (running)"
|
||||
|
||||
|
||||
class TestRemoteResumeSessions:
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_skips_when_vibe_code_disabled(self) -> None:
|
||||
config = MagicMock()
|
||||
config.vibe_code_enabled = False
|
||||
config.vibe_code_api_key = "key"
|
||||
created_clients: list[FakeRemoteResumeClient] = []
|
||||
|
||||
def client_factory(_config: object) -> FakeRemoteResumeClient:
|
||||
client = FakeRemoteResumeClient()
|
||||
created_clients.append(client)
|
||||
return client
|
||||
|
||||
remote = RemoteResumeSessions(lambda: config, client_factory)
|
||||
|
||||
result = await remote.fetch(10.0)
|
||||
|
||||
assert result == RemoteResumeResult([], None)
|
||||
assert created_clients == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_cancels_previous_fetch(self) -> None:
|
||||
config = enabled_vibe_code_config()
|
||||
client = FakeRemoteResumeClient(delay=10.0)
|
||||
remote = RemoteResumeSessions(lambda: config, lambda _config: client)
|
||||
|
||||
first = remote.start(100.0)
|
||||
await asyncio.sleep(0)
|
||||
second = remote.start(100.0)
|
||||
await asyncio.sleep(0)
|
||||
await remote.aclose()
|
||||
|
||||
assert first.cancelled()
|
||||
assert first is not second
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_returns_error_tuple_on_timeout(self) -> None:
|
||||
config = enabled_vibe_code_config()
|
||||
client = FakeRemoteResumeClient(delay=10.0)
|
||||
remote = RemoteResumeSessions(lambda: config, lambda _config: client)
|
||||
|
||||
sessions, error = await remote.fetch(0.01)
|
||||
await remote.aclose()
|
||||
|
||||
assert sessions == []
|
||||
assert error is not None and "Timed out" in error
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aclose_cancels_inflight_fetch_before_closing_client(self) -> None:
|
||||
config = enabled_vibe_code_config()
|
||||
client = FakeRemoteResumeClient(delay=10.0)
|
||||
remote = RemoteResumeSessions(lambda: config, lambda _config: client)
|
||||
|
||||
task = remote.start(100.0)
|
||||
await asyncio.sleep(0)
|
||||
await remote.aclose()
|
||||
|
||||
assert task.cancelled()
|
||||
assert client.closed is True
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue