vibe/tests/setup/test_update_prompt_dialog.py

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, patch

import pytest

from vibe.setup.update_prompt.update_prompt_dialog import (
    UpdateChoice,
    UpdatePromptApp,
    UpdatePromptMode,
    UpdatePromptResult,
)


async def _await_update_completion(app: UpdatePromptApp) -> None:
    while app._update_task is None:
        await asyncio.sleep(0.01)
    try:
        await app._update_task
    except Exception:
        pass


@pytest.mark.asyncio
async def test_dialog_returns_continue_on_continue_selection() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    async with app.run_test() as pilot:
        await pilot.press("right")
        await pilot.press("enter")
        await pilot.pause()

    assert app.return_value is UpdatePromptResult.CONTINUE


@pytest.mark.asyncio
async def test_dialog_returns_updated_when_update_command_succeeds() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    with patch(
        "vibe.setup.update_prompt.update_prompt_dialog.do_update",
        new=AsyncMock(return_value=True),
    ):
        async with app.run_test() as pilot:
            await pilot.press("enter")
            await _await_update_completion(app)
            await pilot.pause()

    assert app.return_value is UpdatePromptResult.UPDATED


@pytest.mark.asyncio
async def test_dialog_returns_update_failed_when_update_command_fails() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    with patch(
        "vibe.setup.update_prompt.update_prompt_dialog.do_update",
        new=AsyncMock(return_value=False),
    ):
        async with app.run_test() as pilot:
            await pilot.press("enter")
            await _await_update_completion(app)
            await pilot.pause()

    assert app.return_value is UpdatePromptResult.UPDATE_FAILED


@pytest.mark.asyncio
async def test_dialog_default_selection_is_update() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    async with app.run_test() as pilot:
        await pilot.pause()
        assert app._dialog is not None
        assert app._dialog.selected is UpdateChoice.UPDATE


@pytest.mark.asyncio
async def test_startup_prompt_uses_continue_label() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    async with app.run_test() as pilot:
        await pilot.pause()
        assert app._dialog is not None
        assert (
            app._dialog._choice_labels[UpdateChoice.CONTINUE]
            == "Continue with current version"
        )


@pytest.mark.asyncio
async def test_check_upgrade_prompt_uses_cancel_label() -> None:
    app = UpdatePromptApp(
        current_version="1.0.0",
        latest_version="2.0.0",
        prompt_mode=UpdatePromptMode.CHECK_UPGRADE,
    )

    async with app.run_test() as pilot:
        await pilot.pause()
        assert app._dialog is not None
        assert app._dialog._choice_labels[UpdateChoice.CONTINUE] == "Cancel upgrade"


@pytest.mark.asyncio
async def test_dialog_returns_quit_on_ctrl_q() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    async with app.run_test() as pilot:
        await pilot.press("ctrl+q")
        await pilot.pause()

    assert app.return_value is UpdatePromptResult.QUIT


@pytest.mark.asyncio
async def test_dialog_returns_update_failed_when_do_update_raises() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    with patch(
        "vibe.setup.update_prompt.update_prompt_dialog.do_update",
        new=AsyncMock(side_effect=OSError("boom")),
    ):
        async with app.run_test() as pilot:
            await pilot.press("enter")
            await _await_update_completion(app)
            await pilot.pause()

    assert app.return_value is UpdatePromptResult.UPDATE_FAILED


@pytest.mark.asyncio
async def test_ctrl_q_during_update_cancels_subprocess_and_quits() -> None:
    app = UpdatePromptApp(current_version="1.0.0", latest_version="2.0.0")

    update_started = asyncio.Event()

    async def slow_update() -> bool:
        update_started.set()
        await asyncio.sleep(60)
        return True

    with patch(
        "vibe.setup.update_prompt.update_prompt_dialog.do_update", new=slow_update
    ):
        async with app.run_test() as pilot:
            await pilot.press("enter")
            await update_started.wait()
            await pilot.press("ctrl+q")
            await pilot.pause()

    assert app.return_value is UpdatePromptResult.QUIT
    assert app._update_task is not None
    assert app._update_task.cancelled() or app._update_task.done()