feat: add Unit.reboot for machine charms

ops/model.py

@@ -691,6 +691,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, after stopping all containers hosted on the machine.
+
+        Normally, the reboot will only take place after the current hook successfully
+        completes. Use the ``now`` argument when multiple reboots are required, to
+        reboot immediately without waiting for the hook to complete, and to restart the
+        hook after rebooting.
+
+        This will silently fail for Kubernetes charms.
+
+        This can only be called for the current unit, and cannot be used in an action
+        hook.
+
+        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 +3383,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

@@ -1882,6 +1882,24 @@ class _Secret:
     grants: Dict[int, Set[str]] = dataclasses.field(default_factory=dict)
 
 
+class RebootingMachineError(Exception):
+    """Raised when the machine would reboot.
+
+    When :meth:`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 a
+    ``RebootingMachineError`` to be raised, for example::
+
+        def test_installation(self):
+            self.harness.begin()
+            with self.assertRaises(ops.testing.RebootMachineError):
+                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:
@@ -2503,6 +2521,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):
+        if not now:
+            # We can't simulate the reboot, so just do nothing.
+            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 RebootingMachineError()
+
 
 @_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,32 @@ 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, *args, **kwargs):  # type: ignore
+                super().__init__(*args, **kwargs)  # type: ignore
+                self.framework.observe(self.on.install, self._reboot)
+                self.framework.observe(self.on.remove, self._reboot_now)
+
+            def _reboot(self, event: ops.RemoveEvent):
+                self.unit.reboot()
+
+            def _reboot_now(self, event: ops.InstallEvent):
+                self.unit.reboot(now=True)
+
+        harness = ops.testing.Harness(RebootingCharm, meta='''
+            name: test-app
+            ''')
+        self.addCleanup(harness.cleanup)
+        backend = harness._backend
+        backend.reboot()
+        with self.assertRaises(ops.testing.RebootingMachineError):
+            backend.reboot(now=True)
+        harness.begin()
+        harness.charm.on.install.emit()
+        with self.assertRaises(ops.testing.RebootingMachineError):
+            harness.charm.on.remove.emit()
+
 
 class _TestingPebbleClientMixin:
     def get_testing_client(self):