From 7065625d2800554aed472e17e6188414562f5071 Mon Sep 17 00:00:00 2001 From: Xitee <59659167+Xitee1@users.noreply.github.com> Date: Mon, 18 Mar 2024 16:20:22 +0100 Subject: [PATCH] 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 | 98 +++++++++++++++++++- tests/components/octoprint/test_button.py | 80 ++++++++++++++++ 2 files changed, 173 insertions(+), 5 deletions(-) diff --git a/homeassistant/components/octoprint/button.py b/homeassistant/components/octoprint/button.py index b7c49df7f4c..2a2e5015303 100644 --- a/homeassistant/components/octoprint/button.py +++ b/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.""" diff --git a/tests/components/octoprint/test_button.py b/tests/components/octoprint/test_button.py index c511227f1b9..39e8fa5886c 100644 --- a/tests/components/octoprint/test_button.py +++ b/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"