Add account selector to Anglian Water config flow (#158242)

This commit is contained in:
Jordan Harvey
2025-12-23 21:04:54 +00:00
committed by GitHub
parent 34db548725
commit b07b699e79
7 changed files with 307 additions and 39 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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