Implement a RebootJob

mqterm/jobs.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from binascii import hexlify
 from hashlib import sha256
@@ -247,6 +248,40 @@ def output(self):
         return BytesIO(str(self.bytes_written).encode("utf-8"))
 
 
+class RebootJob(Job):
+    """A job to perform a hard or soft reboot of the device."""
+
+    argc = 1
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mode = self.args[0]
+
+    def output(self):
+        """Reboot the device after a given delay."""
+        import machine
+
+        msg = f"Performing {self.mode} reboot"
+        try:
+            op = machine.reset if self.mode == "hard" else machine.soft_reset
+        except AttributeError:
+            raise OSError("Operation not supported on this platform")
+        if self.mode == "hard":
+            logging.critical(msg)
+        else:
+            logging.warning(msg)
+
+        # Schedule reboot in three seconds
+        async def reboot_callback(op):
+            await asyncio.sleep(3)
+            op()
+
+        asyncio.create_task(reboot_callback(op))
+
+        # Log the reboot action and return as output
+        return BytesIO(msg.encode("utf-8"))
+
+
 # Map commands to associated job names
 COMMANDS = {
     "whoami": WhoAmIJob,
@@ -255,4 +290,5 @@ def output(self):
     "ls": ListDirJob,
     "cp": PutFileJob,
     "ota": FirmwareUpdateJob,
+    "reboot": RebootJob,
 }

tests/test_jobs.py

@@ -12,6 +12,7 @@
     Job,
     PlatformInfoJob,
     PutFileJob,
+    RebootJob,
     WhoAmIJob,
 )
 
@@ -169,7 +170,7 @@ def test_wait_partial_block(self):
     def test_last_block_fill(self):
         """Should fill space in last block with empty data on final update"""
         initial_data = b"\xde\xad\xbe\xef"
-        checksum = hexlify(sha256(initial_data).digest()).decode('utf-8')
+        checksum = hexlify(sha256(initial_data).digest()).decode("utf-8")
         job = FirmwareUpdateJob("ota", [checksum])
 
         # Send and complete the update
@@ -195,7 +196,7 @@ def test_last_block_fill(self):
     def test_output(self):
         """Should return the total bytes written as output"""
         initial_data = b"\xde\xad\xbe\xef"
-        checksum = hexlify(sha256(initial_data).digest()).decode('utf-8')
+        checksum = hexlify(sha256(initial_data).digest()).decode("utf-8")
         job = FirmwareUpdateJob("ota", [checksum])
 
         # Send and complete the update
@@ -205,3 +206,19 @@ def test_output(self):
         # Output should be the total bytes written
         output = job.output().read().decode("utf-8").strip()
         self.assertEqual(output, "4")
+
+
+class TestRebootJob(TestCase):
+    def test_run_hard(self):
+        """Reboot job should signal and perform a hard reboot"""
+        job = RebootJob("reboot", ["hard"])
+        with self.assertRaises(OSError):  # Can't do on unix
+            output = job.output().read().decode("utf-8").strip()
+            self.assertEqual(output, "Performing hard reboot")
+
+    def test_run_soft(self):
+        """Reboot job should signal a soft reboot"""
+        job = RebootJob("reboot", ["soft"])
+        # TODO: turn off or mock logger here to ignore the output
+        output = job.output().read().decode("utf-8").strip()
+        self.assertEqual(output, "Performing soft reboot")

tests/utils.py

@@ -1,18 +1,21 @@
-# from unittest import TestCase
-
-
 def call(*args, **kwargs):
     return (tuple(args), dict(kwargs))
 
 
 class Mock:
     """A mock callable object that stores its calls."""
 
-    def __init__(self):
+    def __init__(self, return_value=None, side_effect=None):
+        self.return_value = return_value
+        self.side_effect = side_effect
         self._calls = []
 
     def __call__(self, *args, **kwargs):
-        self._calls.append(call(*args, **kwargs))
+        """Call the mock and store the call details."""
+        self._calls.append(call(*args[1:], **kwargs))  # Skip the first argument (self)
+        if self.side_effect:
+            raise self.side_effect
+        return self.return_value
 
     def assert_called(self):
         """Assert that the mock was called at least once."""
@@ -24,17 +27,22 @@ def assert_not_called(self):
 
     def assert_called_with(self, *args, **kwargs):
         """Assert that the mock was last called with the given arguments."""
-        # First call should be self, so we prepend it
-        expected_args = [self] + list(args)
-        expectation = call(*expected_args, **kwargs)
+        # Fail if no calls were made
+        self.assert_called()
 
         # Try to have a useful output for assertion failures
+        expectation = call(*args, **kwargs)
         assert self._calls[-1] == expectation, "Expected call with {}, got {}".format(
             expectation, self._calls[-1]
         )
 
     def assert_has_calls(self, calls):
         """Assert that the mock has the expected calls with arguments."""
+        assert calls, "Expected calls cannot be empty."
+
+        # Fail if no calls were made
+        self.assert_called()
+
         assert self._calls == calls, "Expected calls {}, got {}".format(
             calls, self._calls
         )