Add additional buttons to OctoPrint (#103139)

* Add 3 new buttons

- System shutdown button
- System reboot button
- Octoprint restart button

* Enable buttons by default

* Add tests

* Fix tests

* Remove accidentally committed unused code

* Add RESTART device class to RestartOctoprint and RebootSystem buttons

* Apply suggestions to octoprint test_button

* Freeze time for OctoPrint button tests

* Make new button base class to prevent implementing the availability check multiple times

homeassistant/components/octoprint/button.py

@@ -2,7 +2,7 @@
 
 from pyoctoprintapi import OctoprintClient, OctoprintPrinterInfo
 
-from homeassistant.components.button import ButtonEntity
+from homeassistant.components.button import ButtonDeviceClass, ButtonEntity
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -31,11 +31,16 @@ async def async_setup_entry(
             OctoprintResumeJobButton(coordinator, device_id, client),
             OctoprintPauseJobButton(coordinator, device_id, client),
             OctoprintStopJobButton(coordinator, device_id, client),
+            OctoprintShutdownSystemButton(coordinator, device_id, client),
+            OctoprintRebootSystemButton(coordinator, device_id, client),
+            OctoprintRestartOctoprintButton(coordinator, device_id, client),
         ]
     )
 
 
-class OctoprintButton(CoordinatorEntity[OctoprintDataUpdateCoordinator], ButtonEntity):
+class OctoprintPrinterButton(
+    CoordinatorEntity[OctoprintDataUpdateCoordinator], ButtonEntity
+):
     """Represent an OctoPrint binary sensor."""
 
     client: OctoprintClient
@@ -61,7 +66,35 @@ class OctoprintButton(CoordinatorEntity[OctoprintDataUpdateCoordinator], ButtonE
         return self.coordinator.last_update_success and self.coordinator.data["printer"]
 
 
-class OctoprintPauseJobButton(OctoprintButton):
+class OctoprintSystemButton(
+    CoordinatorEntity[OctoprintDataUpdateCoordinator], ButtonEntity
+):
+    """Represent an OctoPrint binary sensor."""
+
+    client: OctoprintClient
+
+    def __init__(
+        self,
+        coordinator: OctoprintDataUpdateCoordinator,
+        button_type: str,
+        device_id: str,
+        client: OctoprintClient,
+    ) -> None:
+        """Initialize a new OctoPrint button."""
+        super().__init__(coordinator)
+        self.client = client
+        self._device_id = device_id
+        self._attr_name = f"OctoPrint {button_type}"
+        self._attr_unique_id = f"{button_type}-{device_id}"
+        self._attr_device_info = coordinator.device_info
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return self.coordinator.last_update_success
+
+
+class OctoprintPauseJobButton(OctoprintPrinterButton):
     """Pause the active job."""
 
     def __init__(
@@ -83,7 +116,7 @@ class OctoprintPauseJobButton(OctoprintButton):
             raise InvalidPrinterState("Printer is not printing")
 
 
-class OctoprintResumeJobButton(OctoprintButton):
+class OctoprintResumeJobButton(OctoprintPrinterButton):
     """Resume the active job."""
 
     def __init__(
@@ -105,7 +138,7 @@ class OctoprintResumeJobButton(OctoprintButton):
             raise InvalidPrinterState("Printer is not currently paused")
 
 
-class OctoprintStopJobButton(OctoprintButton):
+class OctoprintStopJobButton(OctoprintPrinterButton):
     """Resume the active job."""
 
     def __init__(
@@ -125,5 +158,60 @@ class OctoprintStopJobButton(OctoprintButton):
             await self.client.cancel_job()
 
 
+class OctoprintShutdownSystemButton(OctoprintSystemButton):
+    """Shutdown the system."""
+
+    def __init__(
+        self,
+        coordinator: OctoprintDataUpdateCoordinator,
+        device_id: str,
+        client: OctoprintClient,
+    ) -> None:
+        """Initialize a new OctoPrint button."""
+        super().__init__(coordinator, "Shutdown System", device_id, client)
+
+    async def async_press(self) -> None:
+        """Handle the button press."""
+        await self.client.shutdown()
+
+
+class OctoprintRebootSystemButton(OctoprintSystemButton):
+    """Reboot the system."""
+
+    _attr_device_class = ButtonDeviceClass.RESTART
+
+    def __init__(
+        self,
+        coordinator: OctoprintDataUpdateCoordinator,
+        device_id: str,
+        client: OctoprintClient,
+    ) -> None:
+        """Initialize a new OctoPrint button."""
+        super().__init__(coordinator, "Reboot System", device_id, client)
+
+    async def async_press(self) -> None:
+        """Handle the button press."""
+        await self.client.reboot_system()
+
+
+class OctoprintRestartOctoprintButton(OctoprintSystemButton):
+    """Restart Octoprint."""
+
+    _attr_device_class = ButtonDeviceClass.RESTART
+
+    def __init__(
+        self,
+        coordinator: OctoprintDataUpdateCoordinator,
+        device_id: str,
+        client: OctoprintClient,
+    ) -> None:
+        """Initialize a new OctoPrint button."""
+        super().__init__(coordinator, "Restart Octoprint", device_id, client)
+
+    async def async_press(self) -> None:
+        """Handle the button press."""
+        await self.client.restart()
+
+
 class InvalidPrinterState(HomeAssistantError):
     """Service attempted in invalid state."""

tests/components/octoprint/test_button.py

@@ -2,6 +2,7 @@
 
 from unittest.mock import patch
 
+from freezegun import freeze_time
 from pyoctoprintapi import OctoprintPrinterInfo
 import pytest
 
@@ -193,3 +194,82 @@ async def test_stop_job(hass: HomeAssistant) -> None:
         )
 
         assert len(stop_command.mock_calls) == 0
+
+
+@freeze_time("2023-01-01 00:00")
+async def test_shutdown_system(hass: HomeAssistant) -> None:
+    """Test the shutdown system button."""
+    await init_integration(hass, BUTTON_DOMAIN)
+
+    entity_id = "button.octoprint_shutdown_system"
+
+    # Test shutting down the system
+    with patch(
+        "homeassistant.components.octoprint.coordinator.OctoprintClient.shutdown"
+    ) as shutdown_command:
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+
+        assert len(shutdown_command.mock_calls) == 1
+
+        state = hass.states.get(entity_id)
+        assert state
+        assert state.state == "2023-01-01T00:00:00+00:00"
+
+
+@freeze_time("2023-01-01 00:00")
+async def test_reboot_system(hass: HomeAssistant) -> None:
+    """Test the reboot system button."""
+    await init_integration(hass, BUTTON_DOMAIN)
+
+    entity_id = "button.octoprint_reboot_system"
+
+    # Test rebooting the system
+    with patch(
+        "homeassistant.components.octoprint.coordinator.OctoprintClient.reboot_system"
+    ) as reboot_command:
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {
+                ATTR_ENTITY_ID: entity_id,
+            },
+            blocking=True,
+        )
+
+        assert len(reboot_command.mock_calls) == 1
+
+        state = hass.states.get(entity_id)
+        assert state
+        assert state.state == "2023-01-01T00:00:00+00:00"
+
+
+@freeze_time("2023-01-01 00:00")
+async def test_restart_octoprint(hass: HomeAssistant) -> None:
+    """Test the restart octoprint button."""
+    await init_integration(hass, BUTTON_DOMAIN)
+
+    entity_id = "button.octoprint_restart_octoprint"
+
+    # Test restarting octoprint
+    with patch(
+        "homeassistant.components.octoprint.coordinator.OctoprintClient.restart"
+    ) as restart_command:
+        await hass.services.async_call(
+            BUTTON_DOMAIN,
+            SERVICE_PRESS,
+            {
+                ATTR_ENTITY_ID: entity_id,
+            },
+            blocking=True,
+        )
+
+        assert len(restart_command.mock_calls) == 1
+
+        state = hass.states.get(entity_id)
+        assert state
+        assert state.state == "2023-01-01T00:00:00+00:00"