diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py index abcfe689e93..41e0bd3031a 100644 --- a/homeassistant/components/shelly/config_flow.py +++ b/homeassistant/components/shelly/config_flow.py @@ -58,10 +58,12 @@ async def validate_input( options, ) await rpc_device.shutdown() + assert rpc_device.shelly + return { "title": get_rpc_device_name(rpc_device), CONF_SLEEP_PERIOD: 0, - "model": rpc_device.model, + "model": rpc_device.shelly.get("model"), "gen": 2, } @@ -119,21 +121,21 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): ) except HTTP_CONNECT_ERRORS: errors["base"] = "cannot_connect" - except KeyError: - errors["base"] = "firmware_not_fully_provisioned" except Exception: # pylint: disable=broad-except LOGGER.exception("Unexpected exception") errors["base"] = "unknown" else: - return self.async_create_entry( - title=device_info["title"], - data={ - **user_input, - CONF_SLEEP_PERIOD: device_info[CONF_SLEEP_PERIOD], - "model": device_info["model"], - "gen": device_info["gen"], - }, - ) + if device_info["model"]: + return self.async_create_entry( + title=device_info["title"], + data={ + **user_input, + CONF_SLEEP_PERIOD: device_info[CONF_SLEEP_PERIOD], + "model": device_info["model"], + "gen": device_info["gen"], + }, + ) + errors["base"] = "firmware_not_fully_provisioned" return self.async_show_form( step_id="user", data_schema=HOST_SCHEMA, errors=errors @@ -162,22 +164,22 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): errors["base"] = "cannot_connect" except aioshelly.exceptions.JSONRPCError: errors["base"] = "cannot_connect" - except KeyError: - errors["base"] = "firmware_not_fully_provisioned" except Exception: # pylint: disable=broad-except LOGGER.exception("Unexpected exception") errors["base"] = "unknown" else: - return self.async_create_entry( - title=device_info["title"], - data={ - **user_input, - CONF_HOST: self.host, - CONF_SLEEP_PERIOD: device_info[CONF_SLEEP_PERIOD], - "model": device_info["model"], - "gen": device_info["gen"], - }, - ) + if device_info["model"]: + return self.async_create_entry( + title=device_info["title"], + data={ + **user_input, + CONF_HOST: self.host, + CONF_SLEEP_PERIOD: device_info[CONF_SLEEP_PERIOD], + "model": device_info["model"], + "gen": device_info["gen"], + }, + ) + errors["base"] = "firmware_not_fully_provisioned" else: user_input = {} @@ -223,8 +225,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): try: self.device_info = await validate_input(self.hass, self.host, self.info, {}) - except KeyError: - LOGGER.debug("Shelly host %s firmware not fully provisioned", self.host) except HTTP_CONNECT_ERRORS: return self.async_abort(reason="cannot_connect") @@ -235,7 +235,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): ) -> FlowResult: """Handle discovery confirm.""" errors: dict[str, str] = {} - try: + + if not self.device_info["model"]: + errors["base"] = "firmware_not_fully_provisioned" + model = "Shelly" + else: + model = get_model_name(self.info) if user_input is not None: return self.async_create_entry( title=self.device_info["title"], @@ -246,15 +251,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): "gen": self.device_info["gen"], }, ) - except KeyError: - errors["base"] = "firmware_not_fully_provisioned" - else: self._set_confirm_only() return self.async_show_form( step_id="confirm_discovery", description_placeholders={ - "model": get_model_name(self.info), + "model": model, "host": self.host, }, errors=errors, diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py index 713999de36f..145bcbb3566 100644 --- a/tests/components/shelly/test_config_flow.py +++ b/tests/components/shelly/test_config_flow.py @@ -58,7 +58,7 @@ async def test_form(hass, gen): "aioshelly.rpc_device.RpcDevice.create", new=AsyncMock( return_value=Mock( - model="SHSW-1", + shelly={"model": "SHSW-1", "gen": gen}, config=MOCK_CONFIG, shutdown=AsyncMock(), ) @@ -175,7 +175,7 @@ async def test_form_auth(hass, test_data): "aioshelly.rpc_device.RpcDevice.create", new=AsyncMock( return_value=Mock( - model="SHSW-1", + shelly={"model": "SHSW-1", "gen": gen}, config=MOCK_CONFIG, shutdown=AsyncMock(), ) @@ -225,19 +225,23 @@ async def test_form_errors_get_info(hass, error): assert result2["errors"] == {"base": base_error} -@pytest.mark.parametrize("error", [(KeyError, "firmware_not_fully_provisioned")]) -async def test_form_missing_key_get_info(hass, error): - """Test we handle missing key.""" - exc, base_error = error +async def test_form_missing_model_key(hass): + """Test we handle missing Shelly model key.""" result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_USER} ) with patch( "aioshelly.common.get_info", - return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False, "gen": "2"}, + return_value={"mac": "test-mac", "auth": False, "gen": "2"}, ), patch( - "homeassistant.components.shelly.config_flow.validate_input", - side_effect=KeyError, + "aioshelly.rpc_device.RpcDevice.create", + new=AsyncMock( + return_value=Mock( + shelly={"gen": 2}, + config=MOCK_CONFIG, + shutdown=AsyncMock(), + ) + ), ): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], @@ -245,7 +249,77 @@ async def test_form_missing_key_get_info(hass, error): ) assert result2["type"] == data_entry_flow.RESULT_TYPE_FORM - assert result2["errors"] == {"base": base_error} + assert result2["errors"] == {"base": "firmware_not_fully_provisioned"} + + +async def test_form_missing_model_key_auth_enabled(hass): + """Test we handle missing Shelly model key when auth enabled.""" + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["errors"] == {} + + with patch( + "aioshelly.common.get_info", + return_value={"mac": "test-mac", "auth": True, "gen": 2}, + ): + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"host": "1.1.1.1"}, + ) + + assert result2["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["errors"] == {} + + with patch( + "aioshelly.rpc_device.RpcDevice.create", + new=AsyncMock( + return_value=Mock( + shelly={"gen": 2}, + config=MOCK_CONFIG, + shutdown=AsyncMock(), + ) + ), + ): + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], {"password": "1234"} + ) + + assert result3["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result3["errors"] == {"base": "firmware_not_fully_provisioned"} + + +async def test_form_missing_model_key_zeroconf(hass, caplog): + """Test we handle missing Shelly model key via zeroconf.""" + + with patch( + "aioshelly.common.get_info", + return_value={"mac": "test-mac", "auth": False, "gen": 2}, + ), patch( + "aioshelly.rpc_device.RpcDevice.create", + new=AsyncMock( + return_value=Mock( + shelly={"gen": 2}, + config=MOCK_CONFIG, + shutdown=AsyncMock(), + ) + ), + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=DISCOVERY_INFO, + context={"source": config_entries.SOURCE_ZEROCONF}, + ) + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["errors"] == {"base": "firmware_not_fully_provisioned"} + + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) + assert result2["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result2["errors"] == {"base": "firmware_not_fully_provisioned"} @pytest.mark.parametrize(