vibe/tests/browser_sign_in/stubs.py

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from urllib.parse import urlencode

from vibe.setup.auth import (
    BrowserSignInError,
    BrowserSignInErrorCode,
    BrowserSignInGateway,
    BrowserSignInPollResult,
    BrowserSignInProcess,
)
from vibe.setup.auth.browser_sign_in import BrowserSignInService


@dataclass
class ExchangeRequestPayload:
    process_id: str
    exchange_token: str
    code_verifier: str


class StubBrowserSignInGateway(BrowserSignInGateway):
    def __init__(
        self,
        *,
        process: BrowserSignInProcess | None = None,
        processes: list[BrowserSignInProcess] | None = None,
        poll_results: list[BrowserSignInPollResult | BrowserSignInError] | None = None,
        exchange_result: str = "sk-browser-key",
        exchange_error: BrowserSignInError | None = None,
    ) -> None:
        if process is not None and processes is not None:
            msg = "StubBrowserSignInGateway accepts either process or processes."
            raise AssertionError(msg)

        self._processes = list(processes or ([] if process is None else [process]))
        self._poll_results = list(poll_results or [])
        self.exchange_result = exchange_result
        self.exchange_error = exchange_error
        self.code_challenges: list[str] = []
        self.polled_urls: list[str] = []
        self.exchange_requests: list[ExchangeRequestPayload] = []
        self.closed = False
        self.poll_calls = 0
        self.process_number = 0

    async def create_process(self, code_challenge: str) -> BrowserSignInProcess:
        self.code_challenges.append(code_challenge)
        if not self._processes:
            msg = "StubBrowserSignInGateway requires at least one scripted process."
            raise AssertionError(msg)

        self.process_number += 1
        return self._processes.pop(0)

    async def poll(self, poll_url: str) -> BrowserSignInPollResult:
        self.polled_urls.append(poll_url)
        self.poll_calls += 1
        if not self._poll_results:
            msg = "StubBrowserSignInGateway requires scripted poll results."
            raise AssertionError(msg)
        result = self._poll_results.pop(0)
        if isinstance(result, BrowserSignInError):
            raise result
        return result

    async def exchange(
        self, process_id: str, exchange_token: str, code_verifier: str
    ) -> str:
        self.exchange_requests.append(
            ExchangeRequestPayload(
                process_id=process_id,
                exchange_token=exchange_token,
                code_verifier=code_verifier,
            )
        )
        if self.exchange_error is not None:
            raise self.exchange_error
        return self.exchange_result

    async def aclose(self) -> None:
        self.closed = True


def build_sign_in_process(
    now: datetime, process_id: str = "process-1"
) -> BrowserSignInProcess:
    fragment = urlencode({
        "process_id": process_id,
        "complete_token": f"complete-token-{process_id}",
        "state": f"state-{process_id}",
    })
    return BrowserSignInProcess(
        process_id=process_id,
        sign_in_url=(
            f"https://console.mistral.ai/codestral/cli/authenticate#{fragment}"
        ),
        poll_url=(
            f"https://console.mistral.ai/api/vibe/sign-in/poll/poll-token-{process_id}"
        ),
        expires_at=now + timedelta(minutes=5),
    )


def build_poll_failed_error() -> BrowserSignInError:
    return BrowserSignInError(
        "Browser sign-in status could not be retrieved.",
        code=BrowserSignInErrorCode.POLL_FAILED,
    )


def build_poll_results_from_outcomes(
    outcomes: list[str],
) -> tuple[
    list[BrowserSignInProcess], list[BrowserSignInPollResult | BrowserSignInError]
]:
    processes: list[BrowserSignInProcess] = []
    poll_results: list[BrowserSignInPollResult | BrowserSignInError] = []
    now = datetime(2026, 3, 16, tzinfo=UTC)

    for process_index, outcome in enumerate(outcomes, start=1):
        process_id = f"process-{process_index}"
        processes.append(build_sign_in_process(now, process_id=process_id))

        match outcome:
            case "completed":
                poll_results.extend([
                    BrowserSignInPollResult(status="pending"),
                    BrowserSignInPollResult(
                        status="completed", exchange_token=f"exchange-{process_id}"
                    ),
                ])
            case "expired":
                poll_results.append(BrowserSignInPollResult(status="expired"))
            case "poll_failed":
                poll_results.extend([
                    BrowserSignInPollResult(status="pending"),
                    build_poll_failed_error(),
                    build_poll_failed_error(),
                    build_poll_failed_error(),
                ])
            case _:
                msg = f"Unsupported browser sign-in outcome: {outcome}"
                raise AssertionError(msg)

    return processes, poll_results


async def noop_sleep(_: float) -> None:
    return None


def build_browser_sign_in_service_factory(
    outcomes: list[str],
    *,
    exchange_result: str = "sk-browser-onboarding-test-key",
    open_browser: Callable[[str], bool] | None = None,
    sleep: Callable[[float], Awaitable[None]] = noop_sleep,
    now: Callable[[], datetime] | None = None,
) -> tuple[
    StubBrowserSignInGateway,
    Callable[[], BrowserSignInService],
    list[BrowserSignInService],
]:
    processes, poll_results = build_poll_results_from_outcomes(outcomes)
    gateway = StubBrowserSignInGateway(
        processes=processes, poll_results=poll_results, exchange_result=exchange_result
    )
    created_services: list[BrowserSignInService] = []

    def build_service() -> BrowserSignInService:
        service = BrowserSignInService(
            gateway,
            open_browser=open_browser or (lambda _: True),
            sleep=sleep,
            now=now
            or (
                lambda: (
                    datetime(2026, 3, 16, tzinfo=UTC)
                    + timedelta(seconds=gateway.poll_calls)
                )
            ),
            poll_interval=0,
        )
        created_services.append(service)
        return service

    return gateway, build_service, created_services