Improve code quality of music assistant config flow (home-assistant#156263)

arturpragacz · web-flow · commit c0f61f6c2b40 · 2025-11-11T16:06:54.000+01:00
diff --git a/homeassistant/components/music_assistant/config_flow.py b/homeassistant/components/music_assistant/config_flow.py
@@ -13,29 +13,22 @@
 from music_assistant_models.api import ServerInfoMessage
 import voluptuous as vol
 
-from homeassistant.config_entries import SOURCE_IGNORE, ConfigFlow, ConfigFlowResult
+from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
 from homeassistant.const import CONF_URL
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import aiohttp_client
 from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
 
 from .const import DOMAIN, LOGGER
 
-DEFAULT_URL = "http://mass.local:8095"
 DEFAULT_TITLE = "Music Assistant"
+DEFAULT_URL = "http://mass.local:8095"
 
 
-def get_manual_schema(user_input: dict[str, Any]) -> vol.Schema:
-    """Return a schema for the manual step."""
-    default_url = user_input.get(CONF_URL, DEFAULT_URL)
-    return vol.Schema(
-        {
-            vol.Required(CONF_URL, default=default_url): str,
-        }
-    )
+STEP_USER_SCHEMA = vol.Schema({vol.Required(CONF_URL): str})
 
 
-async def get_server_info(hass: HomeAssistant, url: str) -> ServerInfoMessage:
+async def _get_server_info(hass: HomeAssistant, url: str) -> ServerInfoMessage:
     """Validate the user input allows us to connect."""
     async with MusicAssistantClient(
         url, aiohttp_client.async_get_clientsession(hass)
@@ -52,25 +45,17 @@ class MusicAssistantConfigFlow(ConfigFlow, domain=DOMAIN):
 
     def __init__(self) -> None:
         """Set up flow instance."""
-        self.server_info: ServerInfoMessage | None = None
+        self.url: str | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle a manual configuration."""
         errors: dict[str, str] = {}
+
         if user_input is not None:
             try:
-                self.server_info = await get_server_info(
-                    self.hass, user_input[CONF_URL]
-                )
-                await self.async_set_unique_id(
-                    self.server_info.server_id, raise_on_progress=False
-                )
-                self._abort_if_unique_id_configured(
-                    updates={CONF_URL: user_input[CONF_URL]},
-                    reload_on_update=True,
-                )
+                server_info = await _get_server_info(self.hass, user_input[CONF_URL])
             except CannotConnect:
                 errors["base"] = "cannot_connect"
             except InvalidServerVersion:
@@ -79,85 +64,66 @@ async def async_step_user(
                 LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
+                await self.async_set_unique_id(
+                    server_info.server_id, raise_on_progress=False
+                )
+                self._abort_if_unique_id_configured(
+                    updates={CONF_URL: user_input[CONF_URL]}
+                )
+
                 return self.async_create_entry(
                     title=DEFAULT_TITLE,
-                    data={
-                        CONF_URL: user_input[CONF_URL],
-                    },
+                    data={CONF_URL: user_input[CONF_URL]},
                 )
 
-            return self.async_show_form(
-                step_id="user", data_schema=get_manual_schema(user_input), errors=errors
-            )
+        suggested_values = user_input
+        if suggested_values is None:
+            suggested_values = {CONF_URL: DEFAULT_URL}
 
-        return self.async_show_form(step_id="user", data_schema=get_manual_schema({}))
+        return self.async_show_form(
+            step_id="user",
+            data_schema=self.add_suggested_values_to_schema(
+                STEP_USER_SCHEMA, suggested_values
+            ),
+            errors=errors,
+        )
 
     async def async_step_zeroconf(
         self, discovery_info: ZeroconfServiceInfo
     ) -> ConfigFlowResult:
-        """Handle a discovered Mass server.
-
-        This flow is triggered by the Zeroconf component. It will check if the
-        host is already configured and delegate to the import step if not.
-        """
-        # abort if discovery info is not what we expect
-        if "server_id" not in discovery_info.properties:
-            return self.async_abort(reason="missing_server_id")
-
-        self.server_info = ServerInfoMessage.from_dict(discovery_info.properties)
-        await self.async_set_unique_id(self.server_info.server_id)
+        """Handle a zeroconf discovery for a Music Assistant server."""
+        try:
+            server_info = ServerInfoMessage.from_dict(discovery_info.properties)
+        except LookupError:
+            return self.async_abort(reason="invalid_discovery_info")
 
-        # Check if we already have a config entry for this server_id
-        existing_entry = self.hass.config_entries.async_entry_for_domain_unique_id(
-            DOMAIN, self.server_info.server_id
-        )
+        self.url = server_info.base_url
 
-        if existing_entry:
-            # If the entry was ignored or disabled, don't make any changes
-            if existing_entry.source == SOURCE_IGNORE or existing_entry.disabled_by:
-                return self.async_abort(reason="already_configured")
+        await self.async_set_unique_id(server_info.server_id)
+        self._abort_if_unique_id_configured(updates={CONF_URL: self.url})
 
-            # Test connectivity to the current URL first
-            current_url = existing_entry.data[CONF_URL]
-            try:
-                await get_server_info(self.hass, current_url)
-                # Current URL is working, no need to update
-                return self.async_abort(reason="already_configured")
-            except CannotConnect:
-                # Current URL is not working, update to the discovered URL
-                # and continue to discovery confirm
-                self.hass.config_entries.async_update_entry(
-                    existing_entry,
-                    data={**existing_entry.data, CONF_URL: self.server_info.base_url},
-                )
-                # Schedule reload since URL changed
-                self.hass.config_entries.async_schedule_reload(existing_entry.entry_id)
-        else:
-            # No existing entry, proceed with normal flow
-            self._abort_if_unique_id_configured()
-
-        # Test connectivity to the discovered URL
         try:
-            await get_server_info(self.hass, self.server_info.base_url)
+            await _get_server_info(self.hass, self.url)
         except CannotConnect:
             return self.async_abort(reason="cannot_connect")
+
         return await self.async_step_discovery_confirm()
 
     async def async_step_discovery_confirm(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle user-confirmation of discovered server."""
         if TYPE_CHECKING:
-            assert self.server_info is not None
+            assert self.url is not None
+
         if user_input is not None:
             return self.async_create_entry(
                 title=DEFAULT_TITLE,
-                data={
-                    CONF_URL: self.server_info.base_url,
-                },
+                data={CONF_URL: self.url},
             )
+
         self._set_confirm_only()
         return self.async_show_form(
             step_id="discovery_confirm",
-            description_placeholders={"url": self.server_info.base_url},
+            description_placeholders={"url": self.url},
         )
diff --git a/tests/components/music_assistant/conftest.py b/tests/components/music_assistant/conftest.py
@@ -23,7 +23,7 @@
 def mock_get_server_info() -> Generator[AsyncMock]:
     """Mock the function to get server info."""
     with patch(
-        "homeassistant.components.music_assistant.config_flow.get_server_info"
+        "homeassistant.components.music_assistant.config_flow._get_server_info"
     ) as mock_get_server_info:
         mock_get_server_info.return_value = ServerInfoMessage.from_json(
             load_fixture("server_info_message.json", DOMAIN)
diff --git a/tests/components/music_assistant/test_config_flow.py b/tests/components/music_assistant/test_config_flow.py
@@ -66,7 +66,7 @@ async def test_full_flow(
     assert result["result"].unique_id == "1234"
 
 
-async def test_zero_conf_flow(
+async def test_zeroconf_flow(
     hass: HomeAssistant,
     mock_get_server_info: AsyncMock,
 ) -> None:
@@ -90,21 +90,21 @@ async def test_zero_conf_flow(
     assert result["result"].unique_id == "1234"
 
 
-async def test_zero_conf_missing_server_id(
+async def test_zeroconf_invalid_discovery_info(
     hass: HomeAssistant,
     mock_get_server_info: AsyncMock,
 ) -> None:
-    """Test zeroconf flow with missing server id."""
-    bad_zero_conf_data = deepcopy(ZEROCONF_DATA)
-    bad_zero_conf_data.properties.pop("server_id")
+    """Test zeroconf flow with invalid discovery info."""
+    bad_zeroconf_data = deepcopy(ZEROCONF_DATA)
+    bad_zeroconf_data.properties.pop("server_id")
     result = await hass.config_entries.flow.async_init(
         DOMAIN,
         context={"source": SOURCE_ZEROCONF},
-        data=bad_zero_conf_data,
+        data=bad_zeroconf_data,
     )
     await hass.async_block_till_done()
     assert result["type"] is FlowResultType.ABORT
-    assert result["reason"] == "missing_server_id"
+    assert result["reason"] == "invalid_discovery_info"
 
 
 async def test_duplicate_user(
@@ -324,46 +324,6 @@ async def test_zeroconf_existing_entry_working_url(
     assert mock_config_entry.data[CONF_URL] == "http://localhost:8095"
 
 
-async def test_zeroconf_existing_entry_broken_url(
-    hass: HomeAssistant,
-    mock_get_server_info: AsyncMock,
-    mock_config_entry: MockConfigEntry,
-) -> None:
-    """Test zeroconf flow when existing entry has broken URL."""
-    mock_config_entry.add_to_hass(hass)
-
-    # Create modified zeroconf data with different base_url
-    modified_zeroconf_data = deepcopy(ZEROCONF_DATA)
-    modified_zeroconf_data.properties["base_url"] = "http://discovered-working-url:8095"
-
-    # Mock server info with the discovered URL
-    server_info = ServerInfoMessage.from_json(
-        await async_load_fixture(hass, "server_info_message.json", DOMAIN)
-    )
-    server_info.base_url = "http://discovered-working-url:8095"
-    mock_get_server_info.return_value = server_info
-
-    # First call (testing current URL) should fail, second call (testing discovered URL) should succeed
-    mock_get_server_info.side_effect = [
-        CannotConnect("cannot_connect"),  # Current URL fails
-        server_info,  # Discovered URL works
-    ]
-
-    result = await hass.config_entries.flow.async_init(
-        DOMAIN,
-        context={"source": SOURCE_ZEROCONF},
-        data=modified_zeroconf_data,
-    )
-    await hass.async_block_till_done()
-
-    # Should proceed to discovery confirm because current URL is broken
-    assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == "discovery_confirm"
-    # Verify the URL was updated in the config entry
-    updated_entry = hass.config_entries.async_get_entry(mock_config_entry.entry_id)
-    assert updated_entry.data[CONF_URL] == "http://discovered-working-url:8095"
-
-
 async def test_zeroconf_existing_entry_ignored(
     hass: HomeAssistant,
     mock_get_server_info: AsyncMock,
@@ -396,7 +356,3 @@ async def test_zeroconf_existing_entry_ignored(
     # Should abort because entry was ignored (respect user's choice)
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "already_configured"
-    # Verify the ignored entry was not modified
-    ignored_entry = hass.config_entries.async_get_entry(ignored_config_entry.entry_id)
-    assert ignored_entry.data == {}  # Still no URL field
-    assert ignored_entry.source == SOURCE_IGNORE
