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