feat: add Unit.reboot for machine charms

CHANGES.md

@@ -1,6 +1,6 @@
 # 2.8.0
 
-* Added `Unit.reboot()` and `Harness.last_rebooted``
+* Added `Unit.reboot()` and `Harness.reboot_count``
 * Added `RelationMeta.optional`
 * The type of a `Handle`'s `key` was expanded from `str` to `str|None`
 * Narrowed types of `app` and `unit` in relation events to exclude `None` where applicable

ops/model.py

@@ -26,6 +26,7 @@
 import shutil
 import stat
 import subprocess
+import sys
 import tempfile
 import time
 import typing
@@ -3388,6 +3389,10 @@ def _parse_opened_port(cls, port_str: str) -> Optional[Port]:
     def reboot(self, now: bool = False):
         if now:
             self._run("juju-reboot", "--now")
+            # Juju will kill the Charm process, and in testing no code after
+            # this point would execute. However, we want to guarantee that for
+            # Charmers, so we force that to be the case.
+            sys.exit()
         else:
             self._run("juju-reboot")
 

ops/testing.py

@@ -1733,12 +1733,9 @@ def handle_timeout(args: testing.ExecArgs) -> int:
         )
 
     @property
-    def last_rebooted(self) -> Optional[datetime.datetime]:
-        """The time that the charm last called :meth:`ops.Unit.reboot` (with or without *now*).
-
-        Returns ``None`` if :meth:`ops.Unit.reboot` has not been called.
-        """
-        return self._backend._last_reboot
+    def reboot_count(self) -> int:
+        """Count of times that the charm has called :meth:`ops.Unit.reboot`."""
+        return self._backend._reboot_count
 
 
 def _get_app_or_unit_name(app_or_unit: AppUnitOrName) -> str:
@@ -1974,7 +1971,7 @@ def __init__(self, unit_name: str, meta: charm.CharmMeta, config: '_RawConfig'):
         self._secrets: List[_Secret] = []
         self._opened_ports: Set[model.Port] = set()
         self._networks: Dict[Tuple[Optional[str], Optional[int]], _NetworkDict] = {}
-        self._last_reboot = None
+        self._reboot_count = 0
 
     def _validate_relation_access(self, relation_name: str, relations: List[model.Relation]):
         """Ensures that the named relation exists/has been added.
@@ -2536,11 +2533,14 @@ def _check_protocol_and_port(self, protocol: str, port: Optional[int]):
             raise model.ModelError(f'ERROR invalid protocol "{protocol}", expected "tcp", "udp", or "icmp"\n')  # NOQA: test_quote_backslashes
 
     def reboot(self, now: bool = False):
-        self._last_reboot = datetime.datetime.now(datetime.timezone.utc)
+        self._reboot_count += 1
         if not now:
             return
         # This should exit, reboot, and re-emit the event, but we'll need the caller
         # to handle that. We raise an exception so that they can simulate the exit.
+        # We have a custom exception, rather than using SystemExit like the real call
+        # does, because we don't want to terminate someone's tests if they miss
+        # catching it, or if it happens unexpectedly.
         raise RebootNow()
 
 

test/test_model.py

@@ -3607,7 +3607,8 @@ def test_reboot(self):
         self.assertEqual(fake_script_calls(self, clear=True), [
             ['juju-reboot', ''],
         ])
-        self.unit.reboot(now=True)
+        with self.assertRaises(SystemExit):
+            self.unit.reboot(now=True)
         self.assertEqual(fake_script_calls(self, clear=True), [
             ['juju-reboot', '--now'],
         ])

test/test_testing.py

@@ -3362,27 +3362,19 @@ def _reboot(self, event: ops.RemoveEvent):
             name: test-app
             ''')
         self.addCleanup(harness.cleanup)
-        self.assertIsNone(harness.last_rebooted)
+        self.assertEqual(harness.reboot_count, 0)
         backend = harness._backend
         backend.reboot()
-        reboot1 = harness.last_rebooted
-        assert reboot1 is not None
-        self.assertTrue(reboot1)
+        self.assertEqual(harness.reboot_count, 1)
         with self.assertRaises(ops.testing.RebootNow):
             backend.reboot(now=True)
-        reboot2 = harness.last_rebooted
-        assert reboot2 is not None
-        self.assertGreater(reboot2, reboot1)
+        self.assertEqual(harness.reboot_count, 2)
         harness.begin()
         with self.assertRaises(ops.testing.RebootNow):
             harness.charm.on.install.emit()
-        reboot3 = harness.last_rebooted
-        assert reboot3 is not None
-        self.assertGreater(reboot3, reboot2)
+        self.assertEqual(harness.reboot_count, 3)
         harness.charm.on.remove.emit()
-        reboot4 = harness.last_rebooted
-        assert reboot4 is not None
-        self.assertGreater(reboot4, reboot3)
+        self.assertEqual(harness.reboot_count, 4)
 
 
 class _TestingPebbleClientMixin: