Co-authored-by: Quentin Torroba <quentin.torroba@mistral.ai>
Co-authored-by: Kim-Adeline Miguel <kimadeline.miguel@mistral.ai>
Co-authored-by: Vincent Guilloux <vincent.guilloux@mistral.ai>
Co-authored-by: Clement Sirieix <clem.sirieix@gmail.com>
Co-authored-by: Antoine W <antoine.wronka@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Mathias Gesbert 2026-03-09 19:28:09 +01:00 committed by GitHub
parent 5d2e01a6d7
commit dd372ce494
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
89 changed files with 2086 additions and 596 deletions

View file

@ -8,8 +8,7 @@ import pytest
from tests.cli.plan_offer.adapters.fake_whoami_gateway import FakeWhoAmIGateway
from vibe.cli.plan_offer.decide_plan_offer import (
PlanOfferAction,
PlanType,
WhoAmIPlanType,
decide_plan_offer,
resolve_api_key_for_plan,
)
@ -31,105 +30,96 @@ def mistral_api_key_env() -> Generator[str, None, None]:
@pytest.mark.asyncio
async def test_proposes_upgrade_without_call_when_api_key_is_empty() -> None:
async def test_returns_unknown_plan_when_api_key_is_empty() -> None:
gateway = FakeWhoAmIGateway(
WhoAmIResponse(
is_pro_plan=False,
advertise_pro_plan=False,
plan_type=WhoAmIPlanType.API,
plan_name="Free plan",
prompt_switching_to_pro_plan=False,
)
)
action, plan_type = await decide_plan_offer("", gateway)
plan_info = await decide_plan_offer("", gateway)
assert action is PlanOfferAction.UPGRADE
assert plan_type is PlanType.FREE
assert plan_info.plan_type is WhoAmIPlanType.UNKNOWN
assert plan_info.plan_name == ""
assert plan_info.prompt_switching_to_pro_plan is False
assert gateway.calls == []
@pytest.mark.parametrize(
("response", "expected_action", "expected_plan_type"),
("response", "expected_plan_type", "expected_plan_name", "expected_switch_flag"),
[
(
WhoAmIResponse(
is_pro_plan=True,
advertise_pro_plan=False,
plan_type=WhoAmIPlanType.API,
plan_name="Free Plan",
prompt_switching_to_pro_plan=False,
),
PlanOfferAction.NONE,
PlanType.PRO,
WhoAmIPlanType.API,
"Free Plan",
False,
),
(
WhoAmIResponse(
is_pro_plan=False,
advertise_pro_plan=True,
plan_type=WhoAmIPlanType.CHAT,
plan_name="Pro Plan",
prompt_switching_to_pro_plan=False,
),
PlanOfferAction.UPGRADE,
PlanType.FREE,
WhoAmIPlanType.CHAT,
"Pro Plan",
False,
),
(
WhoAmIResponse(
is_pro_plan=False,
advertise_pro_plan=False,
plan_type=WhoAmIPlanType.CHAT,
plan_name="Pro Plan",
prompt_switching_to_pro_plan=True,
),
PlanOfferAction.SWITCH_TO_PRO_KEY,
PlanType.PRO,
WhoAmIPlanType.CHAT,
"Pro Plan",
True,
),
],
ids=["with-a-pro-plan", "without-a-pro-plan", "with-a-non-pro-key"],
ids=["api-plan", "chat-plan", "chat-plan-with-prompt"],
)
@pytest.mark.asyncio
async def test_proposes_an_action_based_on_current_plan_status(
async def test_returns_plan_info_and_proposes_an_action_based_on_current_plan_status(
response: WhoAmIResponse,
expected_action: PlanOfferAction,
expected_plan_type: PlanType,
expected_plan_type: WhoAmIPlanType,
expected_plan_name: str,
expected_switch_flag: bool,
) -> None:
gateway = FakeWhoAmIGateway(response)
action, plan_type = await decide_plan_offer("api-key", gateway)
plan_info = await decide_plan_offer("api-key", gateway)
assert action is expected_action
assert plan_type is expected_plan_type
assert plan_info.plan_type is expected_plan_type
assert plan_info.plan_name == expected_plan_name
assert plan_info.prompt_switching_to_pro_plan is expected_switch_flag
assert gateway.calls == ["api-key"]
@pytest.mark.asyncio
async def test_proposes_nothing_when_nothing_is_suggested() -> None:
gateway = FakeWhoAmIGateway(
WhoAmIResponse(
is_pro_plan=False,
advertise_pro_plan=False,
prompt_switching_to_pro_plan=False,
)
)
action, plan_type = await decide_plan_offer("api-key", gateway)
assert action is PlanOfferAction.NONE
assert plan_type is PlanType.UNKNOWN
assert gateway.calls == ["api-key"]
@pytest.mark.asyncio
async def test_proposes_upgrade_when_api_key_is_unauthorized() -> None:
async def test_returns_unauthorized_plan_when_api_key_is_unauthorized() -> None:
gateway = FakeWhoAmIGateway(unauthorized=True)
action, plan_type = await decide_plan_offer("bad-key", gateway)
plan_info = await decide_plan_offer("bad-key", gateway)
assert action is PlanOfferAction.UPGRADE
assert plan_type is PlanType.FREE
assert plan_info.plan_type is WhoAmIPlanType.UNAUTHORIZED
assert plan_info.plan_name == ""
assert plan_info.prompt_switching_to_pro_plan is False
assert gateway.calls == ["bad-key"]
@pytest.mark.asyncio
async def test_proposes_none_and_logs_warning_when_gateway_error_occurs(
async def test_returns_unknown_plan_and_logs_warning_when_gateway_error_occurs(
caplog: pytest.LogCaptureFixture,
) -> None:
gateway = FakeWhoAmIGateway(error=True)
with caplog.at_level(logging.WARNING):
action, plan_type = await decide_plan_offer("api-key", gateway)
plan_info = await decide_plan_offer("api-key", gateway)
assert action is PlanOfferAction.NONE
assert plan_type is PlanType.UNKNOWN
assert plan_info.plan_type is WhoAmIPlanType.UNKNOWN
assert plan_info.plan_name == ""
assert plan_info.prompt_switching_to_pro_plan is False
assert gateway.calls == ["api-key"]
assert "Failed to fetch plan status." in caplog.text

View file

@ -8,6 +8,7 @@ from vibe.cli.plan_offer.adapters.http_whoami_gateway import HttpWhoAmIGateway
from vibe.cli.plan_offer.ports.whoami_gateway import (
WhoAmIGatewayError,
WhoAmIGatewayUnauthorized,
WhoAmIPlanType,
WhoAmIResponse,
)
@ -18,8 +19,8 @@ async def test_returns_plan_flags(respx_mock: respx.MockRouter) -> None:
return_value=httpx.Response(
200,
json={
"is_pro_plan": True,
"advertise_pro_plan": False,
"plan_type": "CHAT",
"plan_name": "INDIVIDUAL",
"prompt_switching_to_pro_plan": False,
},
)
@ -31,8 +32,8 @@ async def test_returns_plan_flags(respx_mock: respx.MockRouter) -> None:
assert route.called
request = route.calls.last.request
assert request.headers["Authorization"] == "Bearer api-key"
assert response.is_pro_plan is True
assert response.advertise_pro_plan is False
assert response.plan_type == "CHAT"
assert response.plan_name == "INDIVIDUAL"
assert response.prompt_switching_to_pro_plan is False
@ -68,16 +69,31 @@ async def test_incomplete_payload_defaults_missing_flags_to_false(
respx_mock: respx.MockRouter,
) -> None:
respx_mock.get("http://test/api/vibe/whoami").mock(
return_value=httpx.Response(200, json={"is_pro_plan": True})
return_value=httpx.Response(
200, json={"plan_type": "CHAT", "plan_name": "INDIVIDUAL"}
)
)
gateway = HttpWhoAmIGateway(base_url="http://test")
response = await gateway.whoami("api-key")
assert response == WhoAmIResponse(
is_pro_plan=True, advertise_pro_plan=False, prompt_switching_to_pro_plan=False
plan_type=WhoAmIPlanType.CHAT,
plan_name="INDIVIDUAL",
prompt_switching_to_pro_plan=False,
)
@pytest.mark.asyncio
async def test_raises_on_missing_plan_info(respx_mock: respx.MockRouter) -> None:
respx_mock.get("http://test/api/vibe/whoami").mock(
return_value=httpx.Response(200, json={"prompt_switching_to_pro_plan": False})
)
gateway = HttpWhoAmIGateway(base_url="http://test")
with pytest.raises(WhoAmIGatewayError):
await gateway.whoami("api-key")
@pytest.mark.asyncio
async def test_wraps_request_error(respx_mock: respx.MockRouter) -> None:
respx_mock.get("http://test/api/vibe/whoami").mock(
@ -96,9 +112,9 @@ async def test_parses_boolean_strings(respx_mock: respx.MockRouter) -> None:
return_value=httpx.Response(
200,
json={
"is_pro_plan": "true",
"advertise_pro_plan": "false",
"prompt_switching_to_pro_plan": "true",
"plan_type": "API",
"plan_name": "FREE",
"prompt_switching_to_pro_plan": "false",
},
)
)
@ -106,14 +122,23 @@ async def test_parses_boolean_strings(respx_mock: respx.MockRouter) -> None:
gateway = HttpWhoAmIGateway(base_url="http://test")
response = await gateway.whoami("api-key")
assert response == WhoAmIResponse(
is_pro_plan=True, advertise_pro_plan=False, prompt_switching_to_pro_plan=True
plan_type=WhoAmIPlanType.API,
plan_name="FREE",
prompt_switching_to_pro_plan=False,
)
@pytest.mark.asyncio
async def test_raises_on_invalid_boolean_string(respx_mock: respx.MockRouter) -> None:
respx_mock.get("http://test/api/vibe/whoami").mock(
return_value=httpx.Response(200, json={"is_pro_plan": "yes"})
return_value=httpx.Response(
200,
json={
"plan_type": "CHAT",
"plan_name": "INDIVIDUAL",
"prompt_switching_to_pro_plan": "something",
},
)
)
gateway = HttpWhoAmIGateway(base_url="http://test")