/
/
/
1"""Tests for the Yandex Station 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.ma import BORROW_SOURCE_OWN
13
14from music_assistant.models.setup_flow import SetupFlowContext, SetupSession
15from music_assistant.providers.yandex_station import setup_flow as station_flow
16from music_assistant.providers.yandex_station.constants import (
17 CONF_COOKIES,
18 CONF_MUSIC_TOKEN,
19 CONF_REFRESH_TOKEN,
20 CONF_REMEMBER_SESSION,
21 CONF_X_TOKEN,
22 CONF_YM_INSTANCE,
23)
24
25
26class _FakeClient:
27 """Canned PassportClient that confirms a QR/device login."""
28
29 def __init__(self, creds: Credentials) -> None:
30 self._creds = creds
31
32 async def start_qr_login(self) -> QrSession:
33 return QrSession(track_id="t", csrf_token="c", qr_url="https://passport.yandex.ru/qr/abc")
34
35 async def poll_qr_until_confirmed(self, _qr: QrSession, **_kwargs: Any) -> Credentials:
36 return self._creds
37
38 async def start_device_login(self, **_kwargs: Any) -> DeviceCodeSession:
39 return DeviceCodeSession(
40 device_code=SecretStr("dc"),
41 user_code="ABCD-1234",
42 verification_url="https://ya.ru/device",
43 expires_in=300,
44 interval=5,
45 )
46
47 async def poll_device_until_confirmed(
48 self, _session: DeviceCodeSession, **_kwargs: Any
49 ) -> Credentials:
50 return self._creds
51
52
53def _async_cm(client: _FakeClient) -> mock.MagicMock:
54 """Wrap a fake client as the async context manager PassportClient.create returns."""
55 ctx = mock.MagicMock()
56 ctx.__aenter__ = mock.AsyncMock(return_value=client)
57 ctx.__aexit__ = mock.AsyncMock(return_value=False)
58 return ctx
59
60
61def _make_session(
62 finish_handler: Any, providers: dict[str, Any] | None = None
63) -> tuple[SetupSession, mock.Mock]:
64 """Build a real SetupSession backed by a Mock mass listing the given YM providers."""
65 mass = mock.Mock()
66 mass.config.get = mock.Mock(return_value=providers or {})
67 context = SetupFlowContext(kind="setup", reason="user", domain="yandex_station")
68 return SetupSession(mass, "flow-test", context, finish_handler), mass
69
70
71async def _wait_for(predicate: Any, timeout: float = 5.0) -> Any:
72 """Wait until the predicate returns truthy (or fail the test)."""
73 deadline = time.monotonic() + timeout
74 while time.monotonic() < deadline:
75 if result := predicate():
76 return result
77 await asyncio.sleep(0.01)
78 raise AssertionError("condition not met within timeout")
79
80
81async def _wait_form(session: SetupSession, step_id: str) -> None:
82 """Wait until the given FORM step is presented."""
83 await _wait_for(
84 lambda: (
85 session.current_step
86 and session.current_step.type == FlowStepType.FORM
87 and session.current_step.step_id == step_id
88 )
89 )
90
91
92async def test_borrow_mode_finishes_with_instance_only() -> None:
93 """Selecting a linked Yandex Music instance persists only that instance id."""
94 collected: dict[str, Any] = {}
95
96 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
97 collected.update(values)
98 return {"instance_id": "yandex_station"}
99
100 session, _mass = _make_session(
101 finish, providers={"ym-a": {"domain": "yandex_music", "name": "Main"}}
102 )
103 task = asyncio.create_task(station_flow.run_setup(session))
104 await _wait_form(session, "user")
105 session.handle_submit({CONF_YM_INSTANCE: "ym-a"})
106 await _wait_for(lambda: session.finished)
107 await task
108
109 assert collected == {CONF_YM_INSTANCE: "ym-a"}
110
111
112async def test_own_device_login_persists_full_triple() -> None:
113 """Own credentials via device login persist music + x + refresh tokens under OWN."""
114 creds = Credentials(
115 x_token=SecretStr("XT"),
116 music_token=SecretStr("MT"),
117 refresh_token=SecretStr("RT"),
118 display_login="alice",
119 )
120 collected: dict[str, Any] = {}
121
122 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
123 collected.update(values)
124 return {"instance_id": "yandex_station"}
125
126 session, _mass = _make_session(finish)
127 client = _FakeClient(creds)
128 with mock.patch.object(station_flow, "PassportClient") as pc:
129 pc.create.return_value = _async_cm(client)
130 task = asyncio.create_task(station_flow.run_setup(session))
131 await _wait_form(session, "user")
132 session.handle_submit({CONF_YM_INSTANCE: BORROW_SOURCE_OWN})
133 await _wait_form(session, "method")
134 session.handle_submit(
135 {station_flow.CONF_METHOD: station_flow.METHOD_DEVICE, CONF_REMEMBER_SESSION: True}
136 )
137 await _wait_for(lambda: session.finished)
138 await task
139
140 assert collected == {
141 CONF_YM_INSTANCE: BORROW_SOURCE_OWN,
142 CONF_MUSIC_TOKEN: "MT",
143 CONF_X_TOKEN: "XT",
144 CONF_REFRESH_TOKEN: "RT",
145 }
146
147
148async def test_own_cookie_login_persists_tokens() -> None:
149 """Own credentials via cookies persist music + x token (no refresh) under OWN."""
150 collected: dict[str, Any] = {}
151
152 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
153 collected.update(values)
154 return {"instance_id": "yandex_station"}
155
156 session, _mass = _make_session(finish)
157 with mock.patch.object(
158 station_flow,
159 "login_with_cookies",
160 new=mock.AsyncMock(return_value=("XT", "MT")),
161 ) as cookie_login:
162 task = asyncio.create_task(station_flow.run_setup(session))
163 await _wait_form(session, "user")
164 session.handle_submit({CONF_YM_INSTANCE: BORROW_SOURCE_OWN})
165 await _wait_form(session, "method")
166 session.handle_submit(
167 {station_flow.CONF_METHOD: station_flow.METHOD_COOKIES, CONF_REMEMBER_SESSION: True}
168 )
169 await _wait_form(session, "cookies")
170 session.handle_submit({CONF_COOKIES: "Session_id=abc; yandexuid=1"})
171 await _wait_for(lambda: session.finished)
172 await task
173
174 cookie_login.assert_awaited_once_with("Session_id=abc; yandexuid=1")
175 assert collected == {
176 CONF_YM_INSTANCE: BORROW_SOURCE_OWN,
177 CONF_MUSIC_TOKEN: "MT",
178 CONF_X_TOKEN: "XT",
179 CONF_REFRESH_TOKEN: None,
180 }
181
182
183async def test_own_qr_without_remember_clears_long_lived_tokens() -> None:
184 """QR own login with remember off stores only the music token."""
185 creds = Credentials(x_token=SecretStr("XT"), music_token=SecretStr("MT"))
186 collected: dict[str, Any] = {}
187
188 async def finish(_s: SetupSession, values: dict[str, Any]) -> dict[str, str]:
189 collected.update(values)
190 return {"instance_id": "yandex_station"}
191
192 session, _mass = _make_session(finish)
193 client = _FakeClient(creds)
194 with mock.patch.object(station_flow, "PassportClient") as pc:
195 pc.create.return_value = _async_cm(client)
196 task = asyncio.create_task(station_flow.run_setup(session))
197 await _wait_form(session, "user")
198 session.handle_submit({CONF_YM_INSTANCE: BORROW_SOURCE_OWN})
199 await _wait_form(session, "method")
200 session.handle_submit(
201 {station_flow.CONF_METHOD: station_flow.METHOD_QR, CONF_REMEMBER_SESSION: False}
202 )
203 await _wait_for(lambda: session.finished)
204 await task
205
206 assert collected == {
207 CONF_YM_INSTANCE: BORROW_SOURCE_OWN,
208 CONF_MUSIC_TOKEN: "MT",
209 CONF_X_TOKEN: None,
210 CONF_REFRESH_TOKEN: None,
211 }
212