From a3048234d3750b98046957c710405822be434254 Mon Sep 17 00:00:00 2001 From: Diogo Gomes Date: Tue, 25 Apr 2023 14:16:51 +0100 Subject: [PATCH] Support for multiple contracts in Prosegur (#89097) * support for multiple contracts - tested live * update tests * address review * address review * fix --- .../prosegur/alarm_control_panel.py | 2 +- homeassistant/components/prosegur/camera.py | 4 +- .../components/prosegur/config_flow.py | 61 +++++++++++++------ homeassistant/components/prosegur/const.py | 1 + .../components/prosegur/diagnostics.py | 6 +- .../components/prosegur/manifest.json | 2 +- .../components/prosegur/strings.json | 5 ++ requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/prosegur/conftest.py | 9 +++ tests/components/prosegur/test_config_flow.py | 38 ++++++------ 11 files changed, 87 insertions(+), 45 deletions(-) diff --git a/homeassistant/components/prosegur/alarm_control_panel.py b/homeassistant/components/prosegur/alarm_control_panel.py index cfcb07773f5..b05a5f245ff 100644 --- a/homeassistant/components/prosegur/alarm_control_panel.py +++ b/homeassistant/components/prosegur/alarm_control_panel.py @@ -72,7 +72,7 @@ class ProsegurAlarm(alarm.AlarmControlPanelEntity): """Update alarm status.""" try: - self._installation = await Installation.retrieve(self._auth) + self._installation = await Installation.retrieve(self._auth, self.contract) except ConnectionError as err: _LOGGER.error(err) self._attr_available = False diff --git a/homeassistant/components/prosegur/camera.py b/homeassistant/components/prosegur/camera.py index 848b763903a..9041a6526fb 100644 --- a/homeassistant/components/prosegur/camera.py +++ b/homeassistant/components/prosegur/camera.py @@ -34,7 +34,9 @@ async def async_setup_entry( "async_request_image", ) - _installation = await Installation.retrieve(hass.data[DOMAIN][entry.entry_id]) + _installation = await Installation.retrieve( + hass.data[DOMAIN][entry.entry_id], entry.data["contract"] + ) async_add_entities( [ diff --git a/homeassistant/components/prosegur/config_flow.py b/homeassistant/components/prosegur/config_flow.py index ee2fa795f2d..ea975529b01 100644 --- a/homeassistant/components/prosegur/config_flow.py +++ b/homeassistant/components/prosegur/config_flow.py @@ -11,9 +11,9 @@ from homeassistant import config_entries, core, exceptions from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.data_entry_flow import FlowResult -from homeassistant.helpers import aiohttp_client +from homeassistant.helpers import aiohttp_client, selector -from .const import CONF_COUNTRY, DOMAIN +from .const import CONF_CONTRACT, CONF_COUNTRY, DOMAIN _LOGGER = logging.getLogger(__name__) @@ -31,27 +31,22 @@ async def validate_input(hass: core.HomeAssistant, data): session = aiohttp_client.async_get_clientsession(hass) auth = Auth(session, data[CONF_USERNAME], data[CONF_PASSWORD], data[CONF_COUNTRY]) try: - install = await Installation.retrieve(auth) + contracts = await Installation.list(auth) + return auth, contracts except ConnectionRefusedError: raise InvalidAuth from ConnectionRefusedError except ConnectionError: raise CannotConnect from ConnectionError - # Info to store in the config entry. - return { - "title": f"Contract {install.contract}", - "contract": install.contract, - "username": data[CONF_USERNAME], - "password": data[CONF_PASSWORD], - "country": data[CONF_COUNTRY], - } - class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle a config flow for Prosegur Alarm.""" VERSION = 1 entry: ConfigEntry + auth: Auth + user_input: dict + contracts: list[dict[str, str]] async def async_step_user(self, user_input=None): """Handle the initial step.""" @@ -59,7 +54,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): if user_input: try: - info = await validate_input(self.hass, user_input) + self.auth, self.contracts = await validate_input(self.hass, user_input) except CannotConnect: errors["base"] = "cannot_connect" except InvalidAuth: @@ -68,16 +63,44 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): _LOGGER.exception(exception) errors["base"] = "unknown" else: - await self.async_set_unique_id(info["contract"]) - self._abort_if_unique_id_configured() - - user_input["contract"] = info["contract"] - return self.async_create_entry(title=info["title"], data=user_input) + self.user_input = user_input + return await self.async_step_choose_contract() return self.async_show_form( step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors ) + async def async_step_choose_contract( + self, user_input: Any | None = None + ) -> FlowResult: + """Let user decide which contract is being setup.""" + + if user_input: + await self.async_set_unique_id(user_input[CONF_CONTRACT]) + self._abort_if_unique_id_configured() + + self.user_input[CONF_CONTRACT] = user_input[CONF_CONTRACT] + + return self.async_create_entry( + title=f"Contract {user_input[CONF_CONTRACT]}", data=self.user_input + ) + + contract_options = [ + selector.SelectOptionDict(value=c["contractId"], label=c["description"]) + for c in self.contracts + ] + + return self.async_show_form( + step_id="choose_contract", + data_schema=vol.Schema( + { + vol.Required(CONF_CONTRACT): selector.SelectSelector( + selector.SelectSelectorConfig(options=contract_options) + ), + } + ), + ) + async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult: """Handle initiation of re-authentication with Prosegur.""" self.entry = cast( @@ -93,7 +116,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): if user_input: try: user_input[CONF_COUNTRY] = self.entry.data[CONF_COUNTRY] - await validate_input(self.hass, user_input) + self.auth, self.contracts = await validate_input(self.hass, user_input) except CannotConnect: errors["base"] = "cannot_connect" diff --git a/homeassistant/components/prosegur/const.py b/homeassistant/components/prosegur/const.py index 3f5b8691970..ea823e76062 100644 --- a/homeassistant/components/prosegur/const.py +++ b/homeassistant/components/prosegur/const.py @@ -3,5 +3,6 @@ DOMAIN = "prosegur" CONF_COUNTRY = "country" +CONF_CONTRACT = "contract" SERVICE_REQUEST_IMAGE = "request_image" diff --git a/homeassistant/components/prosegur/diagnostics.py b/homeassistant/components/prosegur/diagnostics.py index d2445698348..59b51f5b5d1 100644 --- a/homeassistant/components/prosegur/diagnostics.py +++ b/homeassistant/components/prosegur/diagnostics.py @@ -9,7 +9,7 @@ from homeassistant.components.diagnostics import async_redact_data from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN +from .const import CONF_CONTRACT, DOMAIN TO_REDACT = {"description", "latitude", "longitude", "contractId", "address"} @@ -19,7 +19,9 @@ async def async_get_config_entry_diagnostics( ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - installation = await Installation.retrieve(hass.data[DOMAIN][entry.entry_id]) + installation = await Installation.retrieve( + hass.data[DOMAIN][entry.entry_id], entry.data[CONF_CONTRACT] + ) activity = await installation.activity(hass.data[DOMAIN][entry.entry_id]) diff --git a/homeassistant/components/prosegur/manifest.json b/homeassistant/components/prosegur/manifest.json index d5081a82dbf..adf5e985fe9 100644 --- a/homeassistant/components/prosegur/manifest.json +++ b/homeassistant/components/prosegur/manifest.json @@ -6,5 +6,5 @@ "documentation": "https://www.home-assistant.io/integrations/prosegur", "iot_class": "cloud_polling", "loggers": ["pyprosegur"], - "requirements": ["pyprosegur==0.0.8"] + "requirements": ["pyprosegur==0.0.9"] } diff --git a/homeassistant/components/prosegur/strings.json b/homeassistant/components/prosegur/strings.json index bf0beb4e766..a6c7fcc4a76 100644 --- a/homeassistant/components/prosegur/strings.json +++ b/homeassistant/components/prosegur/strings.json @@ -8,6 +8,11 @@ "country": "Country" } }, + "choose_contract": { + "data": { + "contract": "Contract" + } + }, "reauth_confirm": { "data": { "description": "Re-authenticate with Prosegur account.", diff --git a/requirements_all.txt b/requirements_all.txt index 70b1686d0b4..ba41b45a41c 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -1886,7 +1886,7 @@ pypoint==2.3.0 pyprof2calltree==1.4.5 # homeassistant.components.prosegur -pyprosegur==0.0.8 +pyprosegur==0.0.9 # homeassistant.components.prusalink pyprusalink==1.1.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 11906ab6acb..94c52b91111 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -1381,7 +1381,7 @@ pypoint==2.3.0 pyprof2calltree==1.4.5 # homeassistant.components.prosegur -pyprosegur==0.0.8 +pyprosegur==0.0.9 # homeassistant.components.prusalink pyprusalink==1.1.0 diff --git a/tests/components/prosegur/conftest.py b/tests/components/prosegur/conftest.py index bd2ce231e28..91bc7f88405 100644 --- a/tests/components/prosegur/conftest.py +++ b/tests/components/prosegur/conftest.py @@ -27,6 +27,15 @@ def mock_config_entry() -> MockConfigEntry: ) +@pytest.fixture +def mock_list_contracts() -> AsyncMock: + """Return list of contracts per user.""" + return [ + {"contractId": "123", "description": "a b c"}, + {"contractId": "456", "description": "x y z"}, + ] + + @pytest.fixture def mock_install() -> AsyncMock: """Return the mocked alarm install.""" diff --git a/tests/components/prosegur/test_config_flow.py b/tests/components/prosegur/test_config_flow.py index 2f57f950a85..2c08f9de109 100644 --- a/tests/components/prosegur/test_config_flow.py +++ b/tests/components/prosegur/test_config_flow.py @@ -1,5 +1,5 @@ """Test the Prosegur Alarm config flow.""" -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -12,7 +12,7 @@ from homeassistant.data_entry_flow import FlowResultType from tests.common import MockConfigEntry -async def test_form(hass: HomeAssistant) -> None: +async def test_form(hass: HomeAssistant, mock_list_contracts) -> None: """Test we get the form.""" result = await hass.config_entries.flow.async_init( @@ -21,12 +21,9 @@ async def test_form(hass: HomeAssistant) -> None: assert result["type"] == "form" assert result["errors"] == {} - install = MagicMock() - install.contract = "123" - with patch( - "homeassistant.components.prosegur.config_flow.Installation.retrieve", - return_value=install, + "homeassistant.components.prosegur.config_flow.Installation.list", + return_value=mock_list_contracts, ) as mock_retrieve, patch( "homeassistant.components.prosegur.async_setup_entry", return_value=True, @@ -41,9 +38,15 @@ async def test_form(hass: HomeAssistant) -> None: ) await hass.async_block_till_done() - assert result2["type"] == "create_entry" - assert result2["title"] == "Contract 123" - assert result2["data"] == { + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], + {"contract": "123"}, + ) + await hass.async_block_till_done() + + assert result3["type"] == "create_entry" + assert result3["title"] == "Contract 123" + assert result3["data"] == { "contract": "123", "username": "test-username", "password": "test-password", @@ -61,7 +64,7 @@ async def test_form_invalid_auth(hass: HomeAssistant) -> None: ) with patch( - "pyprosegur.installation.Installation", + "pyprosegur.installation.Installation.list", side_effect=ConnectionRefusedError, ): result2 = await hass.config_entries.flow.async_configure( @@ -84,7 +87,7 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None: ) with patch( - "pyprosegur.installation.Installation", + "homeassistant.components.prosegur.config_flow.Installation.list", side_effect=ConnectionError, ): result2 = await hass.config_entries.flow.async_configure( @@ -123,7 +126,7 @@ async def test_form_unknown_exception(hass: HomeAssistant) -> None: assert result2["errors"] == {"base": "unknown"} -async def test_reauth_flow(hass: HomeAssistant) -> None: +async def test_reauth_flow(hass: HomeAssistant, mock_list_contracts) -> None: """Test a reauthentication flow.""" entry = MockConfigEntry( domain=DOMAIN, @@ -149,12 +152,9 @@ async def test_reauth_flow(hass: HomeAssistant) -> None: assert result["type"] == FlowResultType.FORM assert result["errors"] == {} - install = MagicMock() - install.contract = "123" - with patch( - "homeassistant.components.prosegur.config_flow.Installation.retrieve", - return_value=install, + "homeassistant.components.prosegur.config_flow.Installation.list", + return_value=mock_list_contracts, ) as mock_installation, patch( "homeassistant.components.prosegur.async_setup_entry", return_value=True, @@ -212,7 +212,7 @@ async def test_reauth_flow_error(hass: HomeAssistant, exception, base_error) -> ) with patch( - "homeassistant.components.prosegur.config_flow.Installation.retrieve", + "homeassistant.components.prosegur.config_flow.Installation.list", side_effect=exception, ): result2 = await hass.config_entries.flow.async_configure(