From c8ac3a29c771546ec38202516ca40902e5f87455 Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Tue, 23 Jan 2024 14:26:47 +0100
Subject: [PATCH] Add reboot and factory_reset to tapodevice (#686)

* Add reboot and factory_reset to tapodevice

* Add test for reboot command

* Fix mocking as different protocols use different methods for comms..
---
 kasa/tapo/tapodevice.py | 15 +++++++++++++++
 kasa/tests/test_cli.py  | 16 ++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/kasa/tapo/tapodevice.py b/kasa/tapo/tapodevice.py
index ff8bdaea..156a61d1 100644
--- a/kasa/tapo/tapodevice.py
+++ b/kasa/tapo/tapodevice.py
@@ -339,3 +339,18 @@ class TapoDevice(SmartDevice):
             "time": t,
         }
         return await self.protocol.query({"set_qs_info": payload})
+
+    async def reboot(self, delay: int = 1) -> None:
+        """Reboot the device.
+
+        Note that giving a delay of zero causes this to block,
+        as the device reboots immediately without responding to the call.
+        """
+        await self.protocol.query({"device_reboot": {"delay": delay}})
+
+    async def factory_reset(self) -> None:
+        """Reset device back to factory settings.
+
+        Note, this does not downgrade the firmware.
+        """
+        await self.protocol.query("device_reset")
diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
index b1db15e1..3aad37dd 100644
--- a/kasa/tests/test_cli.py
+++ b/kasa/tests/test_cli.py
@@ -21,6 +21,7 @@ from kasa.cli import (
     cli,
     emeter,
     raw_command,
+    reboot,
     state,
     sysinfo,
     toggle,
@@ -103,6 +104,21 @@ async def test_raw_command(dev):
     assert "Usage" in res.output
 
 
+@device_smart
+async def test_reboot(dev, mocker):
+    """Test that reboot works on SMART devices."""
+    runner = CliRunner()
+    query_mock = mocker.patch.object(dev.protocol, "query")
+
+    res = await runner.invoke(
+        reboot,
+        obj=dev,
+    )
+
+    query_mock.assert_called()
+    assert res.exit_code == 0
+
+
 @device_smart
 async def test_wifi_scan(dev):
     runner = CliRunner()