From e08661dad370b21ecc1e7cfdfeb7bc23cffcfcc9 Mon Sep 17 00:00:00 2001
From: Brett Adams <Bre77@users.noreply.github.com>
Date: Thu, 31 Aug 2023 17:45:44 +1000
Subject: [PATCH] Patch service validation in Aussie Broadband (#99077)

* Bump pyAussieBB

* rolling back to previous version

* patching the pydantic 2.x issue in aussie_broadband integration

* adding test for validate_service_type

* adding test for validate_service_type

* fixing tests, again

* adding additional test

* doing fixes for live tests

* Implement Feedback

* Add test to detect pydantic2

* Update test_init.py

* Update docstring

---------

Co-authored-by: James Hodgkinson <james@terminaloutcomes.com>
---
 .../components/aussie_broadband/__init__.py   | 19 +++++++++++++++-
 .../components/aussie_broadband/test_init.py  | 22 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/aussie_broadband/__init__.py b/homeassistant/components/aussie_broadband/__init__.py
index ae4bc78580c..1bdb0579976 100644
--- a/homeassistant/components/aussie_broadband/__init__.py
+++ b/homeassistant/components/aussie_broadband/__init__.py
@@ -3,10 +3,11 @@ from __future__ import annotations
 
 from datetime import timedelta
 import logging
+from typing import Any
 
 from aiohttp import ClientError
 from aussiebb.asyncio import AussieBB
-from aussiebb.const import FETCH_TYPES
+from aussiebb.const import FETCH_TYPES, NBN_TYPES, PHONE_TYPES
 from aussiebb.exceptions import AuthenticationException, UnrecognisedServiceType
 
 from homeassistant.config_entries import ConfigEntry
@@ -22,6 +23,19 @@ _LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.SENSOR]
 
 
+# Backport for the pyaussiebb=0.0.15 validate_service_type method
+def validate_service_type(service: dict[str, Any]) -> None:
+    """Check the service types against known types."""
+
+    if "type" not in service:
+        raise ValueError("Field 'type' not found in service data")
+    if service["type"] not in NBN_TYPES + PHONE_TYPES + ["Hardware"]:
+        raise UnrecognisedServiceType(
+            f"Service type {service['type']=} {service['name']=} - not recognised - ",
+            "please report this at https://github.com/yaleman/aussiebb/issues/new",
+        )
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Aussie Broadband from a config entry."""
     # Login to the Aussie Broadband API and retrieve the current service list
@@ -30,6 +44,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         entry.data[CONF_PASSWORD],
         async_get_clientsession(hass),
     )
+    # Overwrite the pyaussiebb=0.0.15 validate_service_type method with backport
+    # Required until pydantic 2.x is supported
+    client.validate_service_type = validate_service_type
     try:
         await client.login()
         services = await client.get_services(drop_types=FETCH_TYPES)
diff --git a/tests/components/aussie_broadband/test_init.py b/tests/components/aussie_broadband/test_init.py
index 3eb1972011c..dc32212ee87 100644
--- a/tests/components/aussie_broadband/test_init.py
+++ b/tests/components/aussie_broadband/test_init.py
@@ -3,8 +3,11 @@ from unittest.mock import patch
 
 from aiohttp import ClientConnectionError
 from aussiebb.exceptions import AuthenticationException, UnrecognisedServiceType
+import pydantic
+import pytest
 
 from homeassistant import data_entry_flow
+from homeassistant.components.aussie_broadband import validate_service_type
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -19,6 +22,19 @@ async def test_unload(hass: HomeAssistant) -> None:
     assert entry.state is ConfigEntryState.NOT_LOADED
 
 
+async def test_validate_service_type() -> None:
+    """Testing the validation function."""
+    test_service = {"type": "Hardware", "name": "test service"}
+    validate_service_type(test_service)
+
+    with pytest.raises(ValueError):
+        test_service = {"name": "test service"}
+        validate_service_type(test_service)
+    with pytest.raises(UnrecognisedServiceType):
+        test_service = {"type": "FunkyBob", "name": "test service"}
+        validate_service_type(test_service)
+
+
 async def test_auth_failure(hass: HomeAssistant) -> None:
     """Test init with an authentication failure."""
     with patch(
@@ -39,3 +55,9 @@ async def test_service_failure(hass: HomeAssistant) -> None:
     """Test init with a invalid service."""
     entry = await setup_platform(hass, usage_effect=UnrecognisedServiceType())
     assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_not_pydantic2() -> None:
+    """Test that Home Assistant still does not support Pydantic 2."""
+    """For PR#99077 and validate_service_type backport"""
+    assert pydantic.__version__ < "2"