From 8f71d7a6f30d0557db474c4d1f6780ada0c73729 Mon Sep 17 00:00:00 2001
From: Duco Sebel <74970928+DCSBL@users.noreply.github.com>
Date: Mon, 13 Jan 2025 13:35:50 +0100
Subject: [PATCH] Move HomeWizard API initialisation to async_setup_entry
 (#135315)

---
 .../components/homewizard/__init__.py         | 21 ++++++++++++++++++-
 .../components/homewizard/coordinator.py      | 14 +++----------
 tests/components/homewizard/conftest.py       |  2 +-
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/homewizard/__init__.py b/homeassistant/components/homewizard/__init__.py
index 4733bc67073..a911f5398da 100644
--- a/homeassistant/components/homewizard/__init__.py
+++ b/homeassistant/components/homewizard/__init__.py
@@ -1,8 +1,12 @@
 """The Homewizard integration."""
 
+from homewizard_energy import HomeWizardEnergy, HomeWizardEnergyV1, HomeWizardEnergyV2
+
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntry
+from homeassistant.const import CONF_IP_ADDRESS, CONF_TOKEN
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
 from .const import DOMAIN, PLATFORMS
 from .coordinator import HWEnergyDeviceUpdateCoordinator
@@ -12,7 +16,22 @@ type HomeWizardConfigEntry = ConfigEntry[HWEnergyDeviceUpdateCoordinator]
 
 async def async_setup_entry(hass: HomeAssistant, entry: HomeWizardConfigEntry) -> bool:
     """Set up Homewizard from a config entry."""
-    coordinator = HWEnergyDeviceUpdateCoordinator(hass)
+
+    api: HomeWizardEnergy
+
+    if token := entry.data.get(CONF_TOKEN):
+        api = HomeWizardEnergyV2(
+            entry.data[CONF_IP_ADDRESS],
+            token=token,
+            clientsession=async_get_clientsession(hass),
+        )
+    else:
+        api = HomeWizardEnergyV1(
+            entry.data[CONF_IP_ADDRESS],
+            clientsession=async_get_clientsession(hass),
+        )
+
+    coordinator = HWEnergyDeviceUpdateCoordinator(hass, api)
     try:
         await coordinator.async_config_entry_first_refresh()
 
diff --git a/homeassistant/components/homewizard/coordinator.py b/homeassistant/components/homewizard/coordinator.py
index f646051a0e1..7024c760b93 100644
--- a/homeassistant/components/homewizard/coordinator.py
+++ b/homeassistant/components/homewizard/coordinator.py
@@ -2,14 +2,12 @@
 
 from __future__ import annotations
 
-from homewizard_energy import HomeWizardEnergy, HomeWizardEnergyV1
+from homewizard_energy import HomeWizardEnergy
 from homewizard_energy.errors import DisabledError, RequestError
 from homewizard_energy.models import CombinedModels as DeviceResponseEntry
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_IP_ADDRESS
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN, LOGGER, UPDATE_INTERVAL
@@ -23,16 +21,10 @@ class HWEnergyDeviceUpdateCoordinator(DataUpdateCoordinator[DeviceResponseEntry]
 
     config_entry: ConfigEntry
 
-    def __init__(
-        self,
-        hass: HomeAssistant,
-    ) -> None:
+    def __init__(self, hass: HomeAssistant, api: HomeWizardEnergy) -> None:
         """Initialize update coordinator."""
         super().__init__(hass, LOGGER, name=DOMAIN, update_interval=UPDATE_INTERVAL)
-        self.api = HomeWizardEnergyV1(
-            self.config_entry.data[CONF_IP_ADDRESS],
-            clientsession=async_get_clientsession(hass),
-        )
+        self.api = api
 
     async def _async_update_data(self) -> DeviceResponseEntry:
         """Fetch all device and sensor data from api."""
diff --git a/tests/components/homewizard/conftest.py b/tests/components/homewizard/conftest.py
index 2e7728c6afb..b540ebac91a 100644
--- a/tests/components/homewizard/conftest.py
+++ b/tests/components/homewizard/conftest.py
@@ -26,7 +26,7 @@ def mock_homewizardenergy(
     """Return a mock bridge."""
     with (
         patch(
-            "homeassistant.components.homewizard.coordinator.HomeWizardEnergyV1",
+            "homeassistant.components.homewizard.HomeWizardEnergyV1",
             autospec=True,
         ) as homewizard,
         patch(