Portainer optimize switch (#163520)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/portainer/switch.py

@@ -41,8 +41,8 @@ class PortainerSwitchEntityDescription(SwitchEntityDescription):
     """Class to hold Portainer switch description."""
 
     is_on_fn: Callable[[PortainerContainerData], bool | None]
-    turn_on_fn: Callable[[str, Portainer, int, str], Coroutine[Any, Any, None]]
-    turn_off_fn: Callable[[str, Portainer, int, str], Coroutine[Any, Any, None]]
+    turn_on_fn: Callable[[Portainer], Callable[[int, str], Coroutine[Any, Any, None]]]
+    turn_off_fn: Callable[[Portainer], Callable[[int, str], Coroutine[Any, Any, None]]]
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -50,53 +50,20 @@ class PortainerStackSwitchEntityDescription(SwitchEntityDescription):
     """Class to hold Portainer stack switch description."""
 
     is_on_fn: Callable[[PortainerStackData], bool | None]
-    turn_on_fn: Callable[[str, Portainer, int, int], Coroutine[Any, Any, None]]
-    turn_off_fn: Callable[[str, Portainer, int, int], Coroutine[Any, Any, None]]
+    turn_on_fn: Callable[[Portainer], Callable[..., Coroutine[Any, Any, Any]]]
+    turn_off_fn: Callable[[Portainer], Callable[..., Coroutine[Any, Any, Any]]]
 
 
 PARALLEL_UPDATES = 1
 
 
-async def perform_container_action(
-    action: str, portainer: Portainer, endpoint_id: int, container_id: str
+async def _perform_action(
+    coordinator: PortainerCoordinator,
+    coroutine: Coroutine[Any, Any, Any],
 ) -> None:
-    """Perform an action on a container."""
+    """Perform a Portainer action with error handling and coordinator refresh."""
     try:
-        match action:
-            case "start":
-                await portainer.start_container(endpoint_id, container_id)
-            case "stop":
-                await portainer.stop_container(endpoint_id, container_id)
-    except PortainerAuthenticationError as err:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="invalid_auth",
-            translation_placeholders={"error": repr(err)},
-        ) from err
-    except PortainerConnectionError as err:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="cannot_connect",
-            translation_placeholders={"error": repr(err)},
-        ) from err
-    except PortainerTimeoutError as err:
-        raise HomeAssistantError(
-            translation_domain=DOMAIN,
-            translation_key="timeout_connect",
-            translation_placeholders={"error": repr(err)},
-        ) from err
-
-
-async def perform_stack_action(
-    action: str, portainer: Portainer, endpoint_id: int, stack_id: int
-) -> None:
-    """Perform an action on a stack."""
-    try:
-        match action:
-            case "start":
-                await portainer.start_stack(stack_id, endpoint_id)
-            case "stop":
-                await portainer.stop_stack(stack_id, endpoint_id)
+        await coroutine
     except PortainerAuthenticationError as err:
         raise HomeAssistantError(
             translation_domain=DOMAIN,
@@ -112,6 +79,8 @@ async def perform_stack_action(
             translation_domain=DOMAIN,
             translation_key="timeout_connect_no_details",
         ) from err
+    else:
+        await coordinator.async_request_refresh()
 
 
 CONTAINER_SWITCHES: tuple[PortainerSwitchEntityDescription, ...] = (
@@ -120,8 +89,8 @@ CONTAINER_SWITCHES: tuple[PortainerSwitchEntityDescription, ...] = (
         translation_key="container",
         device_class=SwitchDeviceClass.SWITCH,
         is_on_fn=lambda data: data.container.state == "running",
-        turn_on_fn=perform_container_action,
-        turn_off_fn=perform_container_action,
+        turn_on_fn=lambda portainer: portainer.start_container,
+        turn_off_fn=lambda portainer: portainer.stop_container,
     ),
 )
 
@@ -131,8 +100,8 @@ STACK_SWITCHES: tuple[PortainerStackSwitchEntityDescription, ...] = (
         translation_key="stack",
         device_class=SwitchDeviceClass.SWITCH,
         is_on_fn=lambda data: data.stack.status == STACK_STATUS_ACTIVE,
-        turn_on_fn=perform_stack_action,
-        turn_off_fn=perform_stack_action,
+        turn_on_fn=lambda portainer: portainer.start_stack,
+        turn_off_fn=lambda portainer: portainer.stop_stack,
     ),
 )
 
@@ -218,23 +187,21 @@ class PortainerContainerSwitch(PortainerContainerEntity, SwitchEntity):
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Start (turn on) the container."""
-        await self.entity_description.turn_on_fn(
-            "start",
-            self.coordinator.portainer,
-            self.endpoint_id,
-            self.container_data.container.id,
+        await _perform_action(
+            self.coordinator,
+            self.entity_description.turn_on_fn(self.coordinator.portainer)(
+                self.endpoint_id, self.container_data.container.id
+            ),
         )
-        await self.coordinator.async_request_refresh()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Stop (turn off) the container."""
-        await self.entity_description.turn_off_fn(
-            "stop",
-            self.coordinator.portainer,
-            self.endpoint_id,
-            self.container_data.container.id,
+        await _perform_action(
+            self.coordinator,
+            self.entity_description.turn_off_fn(self.coordinator.portainer)(
+                self.endpoint_id, self.container_data.container.id
+            ),
         )
-        await self.coordinator.async_request_refresh()
 
 
 class PortainerStackSwitch(PortainerStackEntity, SwitchEntity):
@@ -262,20 +229,18 @@ class PortainerStackSwitch(PortainerStackEntity, SwitchEntity):
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Start (turn on) the stack."""
-        await self.entity_description.turn_on_fn(
-            "start",
-            self.coordinator.portainer,
-            self.endpoint_id,
-            self.stack_data.stack.id,
+        await _perform_action(
+            self.coordinator,
+            self.entity_description.turn_on_fn(self.coordinator.portainer)(
+                self.endpoint_id, self.stack_data.stack.id
+            ),
         )
-        await self.coordinator.async_request_refresh()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Stop (turn off) the stack."""
-        await self.entity_description.turn_off_fn(
-            "stop",
-            self.coordinator.portainer,
-            self.endpoint_id,
-            self.stack_data.stack.id,
+        await _perform_action(
+            self.coordinator,
+            self.entity_description.turn_off_fn(self.coordinator.portainer)(
+                self.endpoint_id, self.stack_data.stack.id
+            ),
         )
-        await self.coordinator.async_request_refresh()