From ac926f9c64a4568118cbb0e45c0e66ab8acb13c4 Mon Sep 17 00:00:00 2001
From: Federico D'Amico <48856240+FedDam@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:55:36 +0100
Subject: [PATCH] Introduce actuator base entity to microBees (#111000)

* changed library imports and added sensors to coordinator and new entity

* revert library

* revert library
---
 .../components/microbees/coordinator.py       |  4 +++-
 homeassistant/components/microbees/entity.py  | 20 +++++++++++++++----
 homeassistant/components/microbees/switch.py  |  4 ++--
 tests/components/microbees/conftest.py        |  4 +++-
 .../components/microbees/test_config_flow.py  |  2 +-
 5 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/microbees/coordinator.py b/homeassistant/components/microbees/coordinator.py
index b085028247c..2e7b5930fd3 100644
--- a/homeassistant/components/microbees/coordinator.py
+++ b/homeassistant/components/microbees/coordinator.py
@@ -7,7 +7,9 @@ from http import HTTPStatus
 import logging
 
 import aiohttp
-from microBeesPy.microbees import Actuator, Bee, MicroBees, MicroBeesException
+from microBeesPy.bee import Actuator, Bee
+from microBeesPy.exceptions import MicroBeesException
+from microBeesPy.microbees import MicroBees
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
diff --git a/homeassistant/components/microbees/entity.py b/homeassistant/components/microbees/entity.py
index e66d7b66f16..5ef90ec60a6 100644
--- a/homeassistant/components/microbees/entity.py
+++ b/homeassistant/components/microbees/entity.py
@@ -1,6 +1,6 @@
 """Base entity for microBees."""
 
-from microBeesPy.microbees import Actuator, Bee
+from microBeesPy.bee import Actuator, Bee
 
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -18,13 +18,10 @@ class MicroBeesEntity(CoordinatorEntity[MicroBeesUpdateCoordinator]):
         self,
         coordinator: MicroBeesUpdateCoordinator,
         bee_id: int,
-        actuator_id: int,
     ) -> None:
         """Initialize the microBees entity."""
         super().__init__(coordinator)
         self.bee_id = bee_id
-        self.actuator_id = actuator_id
-        self._attr_unique_id = f"{bee_id}_{actuator_id}"
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, str(bee_id))},
             manufacturer="microBees",
@@ -46,6 +43,21 @@ class MicroBeesEntity(CoordinatorEntity[MicroBeesUpdateCoordinator]):
         """Return the bee."""
         return self.coordinator.data.bees[self.bee_id]
 
+
+class MicroBeesActuatorEntity(MicroBeesEntity):
+    """Base class for microBees entities with actuator."""
+
+    def __init__(
+        self,
+        coordinator: MicroBeesUpdateCoordinator,
+        bee_id: int,
+        actuator_id: int,
+    ) -> None:
+        """Initialize the microBees entity."""
+        super().__init__(coordinator, bee_id)
+        self.actuator_id = actuator_id
+        self._attr_unique_id = f"{bee_id}_{actuator_id}"
+
     @property
     def actuator(self) -> Actuator:
         """Return the actuator."""
diff --git a/homeassistant/components/microbees/switch.py b/homeassistant/components/microbees/switch.py
index 78a44f44ab3..33db881932e 100644
--- a/homeassistant/components/microbees/switch.py
+++ b/homeassistant/components/microbees/switch.py
@@ -9,7 +9,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
 from .coordinator import MicroBeesUpdateCoordinator
-from .entity import MicroBeesEntity
+from .entity import MicroBeesActuatorEntity
 
 SOCKET_TRANSLATIONS = {46: "socket_it", 38: "socket_eu"}
 
@@ -28,7 +28,7 @@ async def async_setup_entry(
     async_add_entities(switches)
 
 
-class MBSwitch(MicroBeesEntity, SwitchEntity):
+class MBSwitch(MicroBeesActuatorEntity, SwitchEntity):
     """Representation of a microBees switch."""
 
     def __init__(
diff --git a/tests/components/microbees/conftest.py b/tests/components/microbees/conftest.py
index e5886a8dfec..d01fbd2a720 100644
--- a/tests/components/microbees/conftest.py
+++ b/tests/components/microbees/conftest.py
@@ -2,7 +2,9 @@
 import time
 from unittest.mock import AsyncMock, patch
 
-from microBeesPy.microbees import Bee, MicroBees, Profile
+from microBeesPy.bee import Bee
+from microBeesPy.microbees import MicroBees
+from microBeesPy.profile import Profile
 import pytest
 
 from homeassistant.components.application_credentials import (
diff --git a/tests/components/microbees/test_config_flow.py b/tests/components/microbees/test_config_flow.py
index cb93660f743..6962645d6d5 100644
--- a/tests/components/microbees/test_config_flow.py
+++ b/tests/components/microbees/test_config_flow.py
@@ -1,7 +1,7 @@
 """Tests for config flow."""
 from unittest.mock import AsyncMock, patch
 
-from microBeesPy.microbees import MicroBeesException
+from microBeesPy.exceptions import MicroBeesException
 import pytest
 
 from homeassistant.components.microbees.const import DOMAIN