/
/
/
1"""Tests for the Yandex Music interactive setup flow (run_setup)."""
2
3from __future__ import annotations
4
5import asyncio
6import time
7from typing import Any
8from unittest import mock
9
10from music_assistant_models.enums import FlowStepType
11from ya_passport_auth import Credentials, DeviceCodeSession, QrSession, SecretStr
12from ya_passport_auth.exceptions import QRTimeoutError
13
14from music_assistant.models.setup_flow import SetupFlowContext, SetupSession
15from music_assistant.providers.yandex_music import setup_flow as ym_flow
16from music_assistant.providers.yandex_music.constants import (
17 CONF_REFRESH_TOKEN,
18 CONF_REMEMBER_SESSION,
19 CONF_TOKEN,
20 CONF_X_TOKEN,
21)
22
23
24class _FakeClient:
25 """Canned PassportClient that confirms a QR/device login (optionally after one expiry)."""
26
27 def __init__(self, creds: Credentials, *, qr_fail_first: bool = False) -> None:
28 self._creds = creds
29 self.qr_starts = 0
30 self._qr_polls = 0
31 self._qr_fail_first = qr_fail_first
32
33 async def start_qr_login(self) -> QrSession:
34 self.qr_starts += 1
35 return QrSession(track_id="t", csrf_token="c", qr_url="https://passport.yandex.ru/qr/abc")
36
37 async def poll_qr_until_confirmed(self, _qr: QrSession, **_kwargs: Any) -> Credentials:
38 self._qr_polls += 1
39 if self._qr_fail_first and self._qr_polls == 1:
40 raise QRTimeoutError("expired")
41 return self._creds
42
43 async def start_device_login(self, **_kwargs: Any) -> DeviceCodeSession:
44 return DeviceCodeSession(
45 device_code=SecretStr("dc"),
46 user_code="ABCD-1234",
47 verification_url="https://ya.ru/device",
48 expires_in=300,
49 interval=5,
50 )
51
52 async def poll_device_until_confirmed(
53 self, _session: DeviceCodeSession, **_kwargs: Any
54 ) -> Credentials:
55 return self._creds
56
57
58def _async_cm(client: _FakeClient) -> mock.MagicMock:
59 """Wrap a fake client as the async context manager PassportClient.create returns."""
60 ctx = mock.MagicMock()
61 ctx.__aenter__ = mock.AsyncMock(return_value=client)
62 ctx.__aexit__ = mock.AsyncMock(return_value=False)
63 return ctx
64
65
66def _make_session(finish_handler: Any) -> tuple[SetupSession, mock.Mock]:
67 """Build a real SetupSession backed by a Mock mass for driving run_setup directly."""
68 mass = mock.Mock()
69 context = SetupFlowContext(kind="setup", reason="user", domain="yandex_music")
70 return SetupSession(mass, "flow-test", context, finish_handler), mass
71
72
73def _published_steps(mass: mock.Mock) -> list[Any]:
74 """Return the flow steps pushed through mass.signal_event, in order."""
75 return [call.kwargs["data"] for call in mass.signal_event.call_args_list]
76
77
78async def _wait_for(predicate: Any, timeout: float = 5.0) -> Any:
79 """Wait until the predicate returns truthy (or fail the test)."""
80 deadline = time.monotonic() + timeout
81 while time.monotonic() < deadline:
82 if result := predicate():
83 return result
84 await asyncio.sleep(0.01)
85 raise AssertionError("condition not met within timeout")
86
87
88async def _drive(session: SetupSession, submit: dict[str, Any]) -> None:
89 """Wait for the user form, submit the given values, then wait for finish."""
90 await _wait_for(lambda: session.current_step and session.current_step.type == FlowStepType.FORM)
91 session.handle_submit(submit)
92 await _wait_for(lambda: session.finished)
93
94
95async def test_device_login_remember_persists_full_triple() -> None:
96 """Device login with remember on persists music + x + refresh tokens."""
97 creds = Credentials(
98 x_token=SecretStr("XT"),
99 music_token=SecretStr("MT"),
100 refresh_token=SecretStr("RT"),
101 display_login="alice",
102 uid=1,
103 )
104 collected: dict[str, Any] = {}
105
106 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
107 collected.update(values)
108 return {"instance_id": "yandex_music--1"}
109
110 session, mass = _make_session(finish)
111 client = _FakeClient(creds)
112 with mock.patch.object(ym_flow, "PassportClient") as pc:
113 pc.create.return_value = _async_cm(client)
114 task = asyncio.create_task(ym_flow.run_setup(session))
115 await _drive(
116 session, {ym_flow.CONF_METHOD: ym_flow.METHOD_DEVICE, CONF_REMEMBER_SESSION: True}
117 )
118 await task
119
120 assert collected == {CONF_TOKEN: "MT", CONF_X_TOKEN: "XT", CONF_REFRESH_TOKEN: "RT"}
121 progress = [s for s in _published_steps(mass) if s.type == FlowStepType.PROGRESS]
122 assert progress
123 assert progress[0].step_id == "device_login"
124 assert progress[0].image is not None
125 assert progress[0].image.startswith("data:image/svg+xml")
126
127
128async def test_qr_login_without_remember_stores_music_token_only() -> None:
129 """QR login with remember off stores only the music token (x/refresh cleared)."""
130 creds = Credentials(x_token=SecretStr("XT"), music_token=SecretStr("MT"), display_login="bob")
131 collected: dict[str, Any] = {}
132
133 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
134 collected.update(values)
135 return {"instance_id": "yandex_music--1"}
136
137 session, mass = _make_session(finish)
138 client = _FakeClient(creds)
139 with mock.patch.object(ym_flow, "PassportClient") as pc:
140 pc.create.return_value = _async_cm(client)
141 task = asyncio.create_task(ym_flow.run_setup(session))
142 await _drive(
143 session, {ym_flow.CONF_METHOD: ym_flow.METHOD_QR, CONF_REMEMBER_SESSION: False}
144 )
145 await task
146
147 assert collected == {CONF_TOKEN: "MT", CONF_X_TOKEN: None, CONF_REFRESH_TOKEN: None}
148 scan_steps = [s for s in _published_steps(mass) if s.step_id == "scan_qr"]
149 assert scan_steps
150 assert all(s.image and s.image.startswith("data:image/svg+xml") for s in scan_steps)
151
152
153async def test_qr_login_refreshes_expired_code() -> None:
154 """An expired QR code is minted afresh and the login still completes."""
155 creds = Credentials(x_token=SecretStr("XT"), music_token=SecretStr("MT"))
156 collected: dict[str, Any] = {}
157
158 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
159 collected.update(values)
160 return {"instance_id": "yandex_music--1"}
161
162 session, _mass = _make_session(finish)
163 client = _FakeClient(creds, qr_fail_first=True)
164 with mock.patch.object(ym_flow, "PassportClient") as pc:
165 pc.create.return_value = _async_cm(client)
166 task = asyncio.create_task(ym_flow.run_setup(session))
167 await _drive(session, {ym_flow.CONF_METHOD: ym_flow.METHOD_QR, CONF_REMEMBER_SESSION: True})
168 await task
169
170 assert collected[CONF_TOKEN] == "MT"
171 # the expired code triggered a second start_qr_login (refresh loop)
172 assert client.qr_starts == 2
173