Fix Shelly missing key config flow (#72116)

This commit is contained in:
Shay Levy 2022-05-20 12:05:53 +03:00 committed by GitHub
parent 99ad785d0a
commit 74e8b076e5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 117 additions and 41 deletions

View file

@ -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,

View file

@ -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(