From b07b699e79618a8572f7bb93fbc29cd275e752b2 Mon Sep 17 00:00:00 2001 From: Jordan Harvey Date: Tue, 23 Dec 2025 21:04:54 +0000 Subject: [PATCH] Add account selector to Anglian Water config flow (#158242) --- .../components/anglian_water/config_flow.py | 123 ++++++++++++++---- .../components/anglian_water/strings.json | 11 +- tests/components/anglian_water/conftest.py | 20 ++- tests/components/anglian_water/const.py | 2 +- .../fixtures/multi_associated_accounts.json | 65 +++++++++ .../fixtures/single_associated_accounts.json | 37 ++++++ .../anglian_water/test_config_flow.py | 88 +++++++++++-- 7 files changed, 307 insertions(+), 39 deletions(-) create mode 100644 tests/components/anglian_water/fixtures/multi_associated_accounts.json create mode 100644 tests/components/anglian_water/fixtures/single_associated_accounts.json diff --git a/homeassistant/components/anglian_water/config_flow.py b/homeassistant/components/anglian_water/config_flow.py index 5b870d41d3a..ee8fbe4c595 100644 --- a/homeassistant/components/anglian_water/config_flow.py +++ b/homeassistant/components/anglian_water/config_flow.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging -from typing import Any +from typing import TYPE_CHECKING, Any from aiohttp import CookieJar from pyanglianwater import AnglianWater @@ -30,14 +30,11 @@ STEP_USER_DATA_SCHEMA = vol.Schema( vol.Required(CONF_PASSWORD): selector.TextSelector( selector.TextSelectorConfig(type=selector.TextSelectorType.PASSWORD) ), - vol.Required(CONF_ACCOUNT_NUMBER): selector.TextSelector(), } ) -async def validate_credentials( - auth: MSOB2CAuth, account_number: str -) -> str | MSOB2CAuth: +async def validate_credentials(auth: MSOB2CAuth) -> str | MSOB2CAuth: """Validate the provided credentials.""" try: await auth.send_login_request() @@ -46,6 +43,33 @@ async def validate_credentials( except Exception: _LOGGER.exception("Unexpected exception") return "unknown" + return auth + + +def humanize_account_data(account: dict) -> str: + """Convert an account data into a human-readable format.""" + if account["address"]["company_name"] != "": + return f"{account['account_number']} - {account['address']['company_name']}" + if account["address"]["building_name"] != "": + return f"{account['account_number']} - {account['address']['building_name']}" + return f"{account['account_number']} - {account['address']['postcode']}" + + +async def get_accounts(auth: MSOB2CAuth) -> list[selector.SelectOptionDict]: + """Retrieve the list of accounts associated with the authenticated user.""" + _aw = AnglianWater(authenticator=auth) + accounts = await _aw.api.get_associated_accounts() + return [ + selector.SelectOptionDict( + value=str(account["account_number"]), + label=humanize_account_data(account), + ) + for account in accounts["result"]["active"] + ] + + +async def validate_account(auth: MSOB2CAuth, account_number: str) -> str | MSOB2CAuth: + """Validate the provided account number.""" _aw = AnglianWater(authenticator=auth) try: await _aw.validate_smart_meter(account_number) @@ -57,36 +81,91 @@ async def validate_credentials( class AnglianWaterConfigFlow(ConfigFlow, domain=DOMAIN): """Handle a config flow for Anglian Water.""" + def __init__(self) -> None: + """Initialize the config flow.""" + self.authenticator: MSOB2CAuth | None = None + self.accounts: list[selector.SelectOptionDict] = [] + self.user_input: dict[str, Any] | None = None + async def async_step_user( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: """Handle the initial step.""" errors: dict[str, str] = {} if user_input is not None: - validation_response = await validate_credentials( - MSOB2CAuth( - username=user_input[CONF_USERNAME], - password=user_input[CONF_PASSWORD], - session=async_create_clientsession( - self.hass, - cookie_jar=CookieJar(quote_cookie=False), - ), + self.authenticator = MSOB2CAuth( + username=user_input[CONF_USERNAME], + password=user_input[CONF_PASSWORD], + session=async_create_clientsession( + self.hass, + cookie_jar=CookieJar(quote_cookie=False), ), - user_input[CONF_ACCOUNT_NUMBER], ) + validation_response = await validate_credentials(self.authenticator) if isinstance(validation_response, str): errors["base"] = validation_response else: - await self.async_set_unique_id(user_input[CONF_ACCOUNT_NUMBER]) - self._abort_if_unique_id_configured() - return self.async_create_entry( - title=user_input[CONF_ACCOUNT_NUMBER], - data={ - **user_input, - CONF_ACCESS_TOKEN: validation_response.refresh_token, - }, + self.accounts = await get_accounts(self.authenticator) + if len(self.accounts) > 1: + self.user_input = user_input + return await self.async_step_select_account() + account_number = self.accounts[0]["value"] + self.user_input = user_input + return await self.async_step_complete( + { + CONF_ACCOUNT_NUMBER: account_number, + } ) return self.async_show_form( step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors ) + + async def async_step_select_account( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Handle the account selection step.""" + errors = {} + if user_input is not None: + if TYPE_CHECKING: + assert self.authenticator + validation_result = await validate_account( + self.authenticator, + user_input[CONF_ACCOUNT_NUMBER], + ) + if isinstance(validation_result, str): + errors["base"] = validation_result + else: + return await self.async_step_complete(user_input) + return self.async_show_form( + step_id="select_account", + data_schema=vol.Schema( + { + vol.Required(CONF_ACCOUNT_NUMBER): selector.SelectSelector( + selector.SelectSelectorConfig( + options=self.accounts, + multiple=False, + mode=selector.SelectSelectorMode.DROPDOWN, + ) + ) + } + ), + errors=errors, + ) + + async def async_step_complete(self, user_input: dict[str, Any]) -> ConfigFlowResult: + """Handle the final configuration step.""" + await self.async_set_unique_id(user_input[CONF_ACCOUNT_NUMBER]) + self._abort_if_unique_id_configured() + if TYPE_CHECKING: + assert self.authenticator + assert self.user_input + config_entry_data = { + **self.user_input, + CONF_ACCOUNT_NUMBER: user_input[CONF_ACCOUNT_NUMBER], + CONF_ACCESS_TOKEN: self.authenticator.refresh_token, + } + return self.async_create_entry( + title=user_input[CONF_ACCOUNT_NUMBER], + data=config_entry_data, + ) diff --git a/homeassistant/components/anglian_water/strings.json b/homeassistant/components/anglian_water/strings.json index b2c11c1d537..6db91b3b9b0 100644 --- a/homeassistant/components/anglian_water/strings.json +++ b/homeassistant/components/anglian_water/strings.json @@ -10,14 +10,21 @@ "unknown": "[%key:common::config_flow::error::unknown%]" }, "step": { + "select_account": { + "data": { + "account_number": "Billing account number" + }, + "data_description": { + "account_number": "Select the billing account you wish to use." + }, + "description": "Multiple active billing accounts were found with your credentials. Please select the account you wish to use. If this is unexpected, contact Anglian Water to confirm your active accounts." + }, "user": { "data": { - "account_number": "Billing Account Number", "password": "[%key:common::config_flow::data::password%]", "username": "[%key:common::config_flow::data::username%]" }, "data_description": { - "account_number": "Your account number found on your latest bill.", "password": "Your password", "username": "Username or email used to log in to the Anglian Water website." }, diff --git a/tests/components/anglian_water/conftest.py b/tests/components/anglian_water/conftest.py index f206727ad4a..a5106f47791 100644 --- a/tests/components/anglian_water/conftest.py +++ b/tests/components/anglian_water/conftest.py @@ -1,17 +1,19 @@ """Common fixtures for the Anglian Water tests.""" -from collections.abc import Generator +from collections.abc import AsyncGenerator, Generator from unittest.mock import AsyncMock, MagicMock, patch +from pyanglianwater.api import API from pyanglianwater.meter import SmartMeter import pytest from homeassistant.components.anglian_water.const import CONF_ACCOUNT_NUMBER, DOMAIN from homeassistant.const import CONF_ACCESS_TOKEN, CONF_PASSWORD, CONF_USERNAME +from homeassistant.core import HomeAssistant from .const import ACCESS_TOKEN, ACCOUNT_NUMBER, PASSWORD, USERNAME -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, async_load_json_object_fixture @pytest.fixture @@ -67,9 +69,11 @@ def mock_anglian_water_authenticator() -> Generator[MagicMock]: @pytest.fixture -def mock_anglian_water_client( - mock_smart_meter: SmartMeter, mock_anglian_water_authenticator: MagicMock -) -> Generator[AsyncMock]: +async def mock_anglian_water_client( + hass: HomeAssistant, + mock_smart_meter: SmartMeter, + mock_anglian_water_authenticator: MagicMock, +) -> AsyncGenerator[AsyncMock]: """Mock a Anglian Water client.""" # Create a mock instance with our meters and config first. with ( @@ -86,6 +90,12 @@ def mock_anglian_water_client( mock_client.account_config = {"meter_type": "SmartMeter"} mock_client.updated_data_callbacks = [] mock_client.validate_smart_meter.return_value = None + mock_client.api = AsyncMock(spec=API) + mock_client.api.get_associated_accounts.return_value = ( + await async_load_json_object_fixture( + hass, "multi_associated_accounts.json", DOMAIN + ) + ) yield mock_client diff --git a/tests/components/anglian_water/const.py b/tests/components/anglian_water/const.py index 399e7354753..b6a5bbfdb7c 100644 --- a/tests/components/anglian_water/const.py +++ b/tests/components/anglian_water/const.py @@ -1,6 +1,6 @@ """Constants for the Anglian Water test suite.""" -ACCOUNT_NUMBER = "12345678" +ACCOUNT_NUMBER = "171266493" ACCESS_TOKEN = "valid_token" USERNAME = "hello@example.com" PASSWORD = "SecurePassword123" diff --git a/tests/components/anglian_water/fixtures/multi_associated_accounts.json b/tests/components/anglian_water/fixtures/multi_associated_accounts.json new file mode 100644 index 00000000000..2d079c9c1e9 --- /dev/null +++ b/tests/components/anglian_water/fixtures/multi_associated_accounts.json @@ -0,0 +1,65 @@ +{ + "result": { + "property_count": 4, + "active": [ + { + "business_partner_number": 906922831, + "account_number": 171266493, + "address": { + "company_name": "", + "building_name": "", + "sub_building_name": "", + "house_number": "10", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + }, + { + "business_partner_number": 906922832, + "account_number": 171266494, + "address": { + "company_name": "", + "building_name": "Historic Building A", + "sub_building_name": "", + "house_number": "10", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + }, + { + "business_partner_number": 906922832, + "account_number": 171266494, + "address": { + "company_name": "UK Government", + "building_name": "", + "sub_building_name": "", + "house_number": "10", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + } + ], + "inactive": [ + { + "business_partner_number": 100000000, + "account_number": 171200000, + "address": { + "company_name": "", + "building_name": "Finance Office", + "sub_building_name": "Heritage Wing", + "house_number": "50", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + } + ] + } +} diff --git a/tests/components/anglian_water/fixtures/single_associated_accounts.json b/tests/components/anglian_water/fixtures/single_associated_accounts.json new file mode 100644 index 00000000000..04c2033767f --- /dev/null +++ b/tests/components/anglian_water/fixtures/single_associated_accounts.json @@ -0,0 +1,37 @@ +{ + "result": { + "property_count": 1, + "active": [ + { + "business_partner_number": 906922831, + "account_number": 171266493, + "address": { + "company_name": "", + "building_name": "", + "sub_building_name": "", + "house_number": "10", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + } + ], + "inactive": [ + { + "business_partner_number": 100000000, + "account_number": 171200000, + "address": { + "company_name": "", + "building_name": "Finance Office", + "sub_building_name": "Heritage Wing", + "house_number": "50", + "street": "DOWNING STREET", + "locality": "", + "city": "LONDON", + "postcode": "SW1A 1AA" + } + } + ] + } +} diff --git a/tests/components/anglian_water/test_config_flow.py b/tests/components/anglian_water/test_config_flow.py index d577a35880a..8ce8710bf8b 100644 --- a/tests/components/anglian_water/test_config_flow.py +++ b/tests/components/anglian_water/test_config_flow.py @@ -18,16 +18,16 @@ from homeassistant.data_entry_flow import FlowResultType from .const import ACCESS_TOKEN, ACCOUNT_NUMBER, PASSWORD, USERNAME -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, async_load_json_object_fixture -async def test_full_flow( +async def test_multiple_account_flow( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_anglian_water_authenticator: AsyncMock, mock_anglian_water_client: AsyncMock, ) -> None: - """Test a full and successful config flow.""" + """Test the config flow when there are multiple accounts.""" result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER} ) @@ -40,6 +40,15 @@ async def test_full_flow( user_input={ CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD, + }, + ) + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "select_account" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, ) @@ -53,6 +62,43 @@ async def test_full_flow( assert result["result"].unique_id == ACCOUNT_NUMBER +async def test_single_account_flow( + hass: HomeAssistant, + mock_setup_entry: AsyncMock, + mock_anglian_water_authenticator: AsyncMock, + mock_anglian_water_client: AsyncMock, +) -> None: + """Test the config flow when there is just a single account.""" + mock_anglian_water_client.api.get_associated_accounts.return_value = ( + await async_load_json_object_fixture( + hass, "single_associated_accounts.json", DOMAIN + ) + ) + + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": SOURCE_USER} + ) + assert result is not None + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "user" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_USERNAME: USERNAME, + CONF_PASSWORD: PASSWORD, + }, + ) + + assert result["type"] is FlowResultType.CREATE_ENTRY + assert result["title"] == ACCOUNT_NUMBER + assert result["data"][CONF_USERNAME] == USERNAME + assert result["data"][CONF_PASSWORD] == PASSWORD + assert result["data"][CONF_ACCESS_TOKEN] == ACCESS_TOKEN + assert result["data"][CONF_ACCOUNT_NUMBER] == ACCOUNT_NUMBER + assert result["result"].unique_id == ACCOUNT_NUMBER + + async def test_already_configured( hass: HomeAssistant, mock_setup_entry: AsyncMock, @@ -75,6 +121,15 @@ async def test_already_configured( user_input={ CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD, + }, + ) + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "select_account" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, ) @@ -109,7 +164,6 @@ async def test_auth_recover_exception( user_input={ CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD, - CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, ) @@ -126,6 +180,15 @@ async def test_auth_recover_exception( user_input={ CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD, + }, + ) + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "select_account" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, ) @@ -161,19 +224,28 @@ async def test_account_recover_exception( assert result["type"] is FlowResultType.FORM assert result["step_id"] == "user" - mock_anglian_water_client.validate_smart_meter.side_effect = exception_type - result = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={ CONF_USERNAME: USERNAME, CONF_PASSWORD: PASSWORD, + }, + ) + + mock_anglian_water_client.validate_smart_meter.side_effect = exception_type + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "select_account" + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, ) assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "user" + assert result["step_id"] == "select_account" assert result["errors"] == {"base": expected_error} # Now test we can recover @@ -183,8 +255,6 @@ async def test_account_recover_exception( result = await hass.config_entries.flow.async_configure( result["flow_id"], user_input={ - CONF_USERNAME: USERNAME, - CONF_PASSWORD: PASSWORD, CONF_ACCOUNT_NUMBER: ACCOUNT_NUMBER, }, )