Support for multiple contracts in Prosegur (#89097)

* support for multiple contracts - tested live

* update tests

* address review

* address review

* fix

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

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

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.user_input = user_input
-                self._abort_if_unique_id_configured()
+                return await self.async_step_choose_contract()
-
-                user_input["contract"] = info["contract"]
-                return self.async_create_entry(title=info["title"], data=user_input)
 
         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"

homeassistant/components/prosegur/const.py

@@ -3,5 +3,6 @@
 DOMAIN = "prosegur"
 
 CONF_COUNTRY = "country"
+CONF_CONTRACT = "contract"
 
 SERVICE_REQUEST_IMAGE = "request_image"

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])
 

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"]
 }

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.",

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

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

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."""

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",
+        "homeassistant.components.prosegur.config_flow.Installation.list",
-        return_value=install,
+        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"
+        result3 = await hass.config_entries.flow.async_configure(
-    assert result2["title"] == "Contract 123"
+            result2["flow_id"],
-    assert result2["data"] == {
+            {"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",
+        "homeassistant.components.prosegur.config_flow.Installation.list",
-        return_value=install,
+        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(