/
/
/
1"""Tests for the AI Radio setup flow."""
2
3from __future__ import annotations
4
5import asyncio
6import time
7from types import SimpleNamespace
8from typing import Any
9from unittest.mock import AsyncMock, MagicMock
10
11import pytest
12from music_assistant_models.enums import FlowStepType, ProviderFeature
13
14from music_assistant.models.plugin import AIEngine, PluginProvider, TTSEngine
15from music_assistant.models.setup_flow import (
16 AbortFlow,
17 SetupFlowContext,
18 SetupSession,
19)
20from music_assistant.providers.ai_radio import setup_flow
21
22
23def _create_plugin(instance_id: str, ai_ids: list[str], tts_ids: list[str]) -> MagicMock:
24 """Create a mock plugin provider exposing the given AI and TTS engines."""
25 provider = MagicMock(spec=PluginProvider)
26 provider.instance_id = instance_id
27 provider.get_ai_engines = AsyncMock(
28 return_value=[
29 AIEngine(id=engine_id, name=engine_id, provider=provider) for engine_id in ai_ids
30 ]
31 )
32 provider.get_tts_engines = AsyncMock(
33 return_value=[
34 TTSEngine(id=engine_id, name=engine_id, provider=provider) for engine_id in tts_ids
35 ]
36 )
37 return provider
38
39
40def _create_mass(*plugins: MagicMock, locale: str = "en_US") -> MagicMock:
41 """Create a mock MusicAssistant serving the given plugins for the engine features."""
42 mass = MagicMock()
43 mass.metadata = SimpleNamespace(locale=locale)
44
45 def _providers(feature: ProviderFeature, priority: Any = ()) -> list[MagicMock]: # noqa: ARG001
46 if feature == ProviderFeature.AI_QUERY:
47 return [plugin for plugin in plugins if plugin.get_ai_engines.return_value]
48 return [plugin for plugin in plugins if plugin.get_tts_engines.return_value]
49
50 mass.get_providers_supporting_feature.side_effect = _providers
51 return mass
52
53
54def _create_session(
55 mass: MagicMock,
56 collected: dict[str, Any],
57 setup_data: dict[str, Any] | None = None,
58) -> SetupSession:
59 """Create a setup session recording the values handed to finish()."""
60
61 async def finish(_session: SetupSession, values: dict[str, Any]) -> dict[str, str]:
62 collected.update(values)
63 return {"instance_id": "ai_radio--1"}
64
65 return SetupSession(
66 mass,
67 "flow1",
68 SetupFlowContext(
69 kind="setup",
70 reason="user",
71 domain="ai_radio",
72 setup_data=setup_data or {},
73 ),
74 finish,
75 )
76
77
78async def _wait_for_form(session: SetupSession) -> Any:
79 """Wait until the flow published its form step."""
80 deadline = time.monotonic() + 5
81 while time.monotonic() < deadline:
82 step = session.current_step
83 if step is not None and step.type == FlowStepType.FORM:
84 return step
85 await asyncio.sleep(0.01)
86 raise AssertionError("setup flow did not publish a form step")
87
88
89async def test_setup_collects_both_engines() -> None:
90 """The single form offers an AI and a TTS picker and persists the selection."""
91 mass = _create_mass(_create_plugin("hass--1", ["ai_task.a"], ["tts.a"]))
92 collected: dict[str, Any] = {}
93 session = _create_session(mass, collected)
94
95 task = asyncio.create_task(setup_flow.run_setup(session))
96 step = await _wait_for_form(session)
97 assert [entry.key for entry in step.entries] == [
98 "ai_engine",
99 "tts_engine",
100 "weather_city",
101 "weather_country",
102 ]
103 assert [option.value for option in step.entries[0].options] == ["hass--1/ai_task.a"]
104 assert (
105 session.handle_submit({"ai_engine": "hass--1/ai_task.a", "tts_engine": "hass--1/tts.a"})
106 is None
107 )
108 await task
109
110 assert collected == {
111 "ai_engine": "hass--1/ai_task.a",
112 "tts_engine": "hass--1/tts.a",
113 "weather_city": "",
114 "weather_country": "US",
115 }
116
117
118async def test_setup_offers_the_weather_location() -> None:
119 """A city/country submitted at setup is persisted alongside the engine picks."""
120 mass = _create_mass(_create_plugin("hass--1", ["ai_task.a"], ["tts.a"]))
121 collected: dict[str, Any] = {}
122 session = _create_session(mass, collected)
123
124 task = asyncio.create_task(setup_flow.run_setup(session))
125 await _wait_for_form(session)
126 session.handle_submit(
127 {
128 "ai_engine": "hass--1/ai_task.a",
129 "tts_engine": "hass--1/tts.a",
130 "weather_city": "Amsterdam",
131 "weather_country": "NL",
132 }
133 )
134 await task
135
136 assert collected["weather_city"] == "Amsterdam"
137 assert collected["weather_country"] == "NL"
138
139
140async def test_setup_rejects_a_submission_without_a_choice() -> None:
141 """Both engines are mandatory, so an empty submission re-shows the form with an error."""
142 mass = _create_mass(_create_plugin("hass--1", ["ai_task.a"], ["tts.a"]))
143 collected: dict[str, Any] = {}
144 session = _create_session(mass, collected)
145
146 task = asyncio.create_task(setup_flow.run_setup(session))
147 await _wait_for_form(session)
148 step = session.handle_submit({})
149 assert step is not None
150 assert step.errors
151 assert not collected
152
153 session.handle_submit({"ai_engine": "hass--1/ai_task.a", "tts_engine": "hass--1/tts.a"})
154 await task
155
156
157async def test_setup_prefills_the_previous_selection() -> None:
158 """A reconfigure run starts from the values collected by the previous run."""
159 mass = _create_mass(_create_plugin("hass--1", ["ai_task.a"], ["tts.a"]))
160 session = _create_session(mass, {}, setup_data={"tts_engine": "hass--1/tts.a"})
161
162 task = asyncio.create_task(setup_flow.run_setup(session))
163 step = await _wait_for_form(session)
164 assert step.entries[0].value is None
165 assert step.entries[1].value == "hass--1/tts.a"
166 session.handle_submit({"ai_engine": "hass--1/ai_task.a", "tts_engine": "hass--1/tts.a"})
167 await task
168
169
170async def test_setup_aborts_without_an_ai_engine() -> None:
171 """No AI engine means the user has to set up a plugin providing AI first."""
172 mass = _create_mass(_create_plugin("hass--1", [], ["tts.a"]))
173
174 with pytest.raises(AbortFlow) as error:
175 await setup_flow.run_setup(_create_session(mass, {}))
176
177 assert error.value.reason == "no_ai_engine"
178
179
180async def test_setup_aborts_without_a_tts_engine() -> None:
181 """No TTS engine means the user has to set up a plugin providing speech first."""
182 mass = _create_mass(_create_plugin("hass--1", ["ai_task.a"], []))
183
184 with pytest.raises(AbortFlow) as error:
185 await setup_flow.run_setup(_create_session(mass, {}))
186
187 assert error.value.reason == "no_tts_engine"
188