feat: add Unit.reboot for machine charms

CHANGES.md

@@ -1,6 +1,9 @@
 # 2.8.0
 
+* Added `Unit.reboot()` and `Harness.last_rebooted``
+* 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
 
 # 2.7.0
 

ops/model.py

@@ -673,6 +673,8 @@ def set_ports(self, *ports: Union[int, 'Port']) -> None:
         Use :meth:`open_port` and :meth:`close_port` to manage ports
         individually.
 
+        *New in version 2.7*
+
         Args:
             ports: The ports to open. Provide an int to open a TCP port, or
                 a :class:`Port` to open a port for another protocol.
@@ -691,6 +693,32 @@ def set_ports(self, *ports: Union[int, 'Port']) -> None:
         for protocol, port in desired - existing:
             self._backend.open_port(protocol, port)
 
+    def reboot(self, now: bool = False) -> None:
+        """Reboot the host machine.
+
+        Normally, the reboot will only take place after the current hook successfully
+        completes. Use ``now=True`` to reboot immediately without waiting for the
+        hook to complete; this is useful when multiple restarts are required (Juju
+        will re-run the hook after rebooting).
+
+        This is not supported on Kubernetes charms, can only be called for the current unit,
+        and cannot be used in an action hook.
+
+        *New in version 2.8*
+
+        Args:
+            now: terminate immediately without waiting for the current hook to complete,
+                restarting the hook after reboot.
+
+        Raises:
+            RuntimeError: if called on a remote unit.
+            :class:`ModelError`: if used in an action hook.
+
+        """
+        if not self._is_our_unit:
+            raise RuntimeError(f'cannot reboot a remote unit {self}')
+        self._backend.reboot(now)
+
 
 @dataclasses.dataclass(frozen=True)
 class Port:
@@ -3357,6 +3385,12 @@ def _parse_opened_port(cls, port_str: str) -> Optional[Port]:
         protocol_lit = typing.cast(typing.Literal['tcp', 'udp'], protocol)
         return Port(protocol_lit, int(port))
 
+    def reboot(self, now: bool = False):
+        if now:
+            self._run("juju-reboot", "--now")
+        else:
+            self._run("juju-reboot")
+
 
 class _ModelBackendValidator:
     """Provides facilities for validating inputs and formatting them for model backends."""

ops/testing.py

@@ -1727,6 +1727,14 @@ def handle_timeout(args: testing.ExecArgs) -> int:
             handler=(lambda _: result) if handler is None else handler  # type: ignore
         )
 
+    @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 _get_app_or_unit_name(app_or_unit: AppUnitOrName) -> str:
     """Return name of given application or unit (return strings directly)."""
@@ -1882,6 +1890,24 @@ class _Secret:
     grants: Dict[int, Set[str]] = dataclasses.field(default_factory=dict)
 
 
+class RebootNow(Exception):  # noqa
+    """Raised when the machine would reboot.
+
+    When :meth:`ops.Unit.reboot` is called with ``now=True`` in a machine charm, the
+    unit's machine is rebooted, interrupting the execution of the event hook. To
+    simulate that when using the testing harness, write a test that expects this
+    exception to be raised, for example::
+
+        def test_installation(self):
+            self.harness.begin()
+            with self.assertRaises(ops.testing.RebootNow):
+                self.harness.charm.on.install.emit()
+                # More asserts here that the first part of installation was done.
+            self.harness.charm.on.install.emit()
+            # More asserts here that the installation continued appropriately.
+    """
+
+
 @_copy_docstrings(model._ModelBackend)
 @_record_calls
 class _TestingModelBackend:
@@ -1943,6 +1969,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
 
     def _validate_relation_access(self, relation_name: str, relations: List[model.Relation]):
         """Ensures that the named relation exists/has been added.
@@ -2503,6 +2530,14 @@ def _check_protocol_and_port(self, protocol: str, port: Optional[int]):
         else:
             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)
+        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.
+        raise RebootNow()
+
 
 @_copy_docstrings(pebble.ExecProcess)
 class _TestingExecProcess:

test/test_model.py

@@ -3596,5 +3596,27 @@ def test_set_ports_noop(self):
         ])
 
 
+class TestUnit(unittest.TestCase):
+    def setUp(self):
+        self.model = ops.model.Model(ops.charm.CharmMeta(), ops.model._ModelBackend('myapp/0'))
+        self.unit = self.model.unit
+
+    def test_reboot(self):
+        fake_script(self, 'juju-reboot', 'exit 0')
+        self.unit.reboot()
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['juju-reboot', ''],
+        ])
+        self.unit.reboot(now=True)
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['juju-reboot', '--now'],
+        ])
+
+        with self.assertRaises(RuntimeError):
+            self.model.get_unit('other').reboot()
+        with self.assertRaises(RuntimeError):
+            self.model.get_unit('other').reboot(now=True)
+
+
 if __name__ == "__main__":
     unittest.main()

test/test_testing.py

@@ -3337,6 +3337,45 @@ def test_get_pebble_methods(self):
         client = backend.get_pebble('/custom/socket/path')
         self.assertIsInstance(client, _TestingPebbleClient)
 
+    def test_reboot(self):
+        class RebootingCharm(ops.CharmBase):
+            def __init__(self, framework: ops.Framework):
+                super().__init__(framework)
+                self.framework.observe(self.on.install, self._reboot_now)
+                self.framework.observe(self.on.remove, self._reboot)
+
+            def _reboot_now(self, event: ops.InstallEvent):
+                self.unit.reboot(now=True)
+
+            def _reboot(self, event: ops.RemoveEvent):
+                self.unit.reboot()
+
+        harness = ops.testing.Harness(RebootingCharm, meta='''
+            name: test-app
+            ''')
+        self.addCleanup(harness.cleanup)
+        self.assertIsNone(harness.last_rebooted)
+        backend = harness._backend
+        backend.reboot()
+        reboot1 = harness.last_rebooted
+        assert reboot1 is not None
+        self.assertTrue(reboot1)
+        with self.assertRaises(ops.testing.RebootNow):
+            backend.reboot(now=True)
+        reboot2 = harness.last_rebooted
+        assert reboot2 is not None
+        self.assertGreater(reboot2, reboot1)
+        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)
+        harness.charm.on.remove.emit()
+        reboot4 = harness.last_rebooted
+        assert reboot4 is not None
+        self.assertGreater(reboot4, reboot3)
+
 
 class _TestingPebbleClientMixin:
     def get_testing_client(self):