feat: add Unit.reboot for machine charms (#1041)

Adds a new `reboot` method to `Unit`, that's a wrapper around the `juju-reboot` command.

`juju-reboot` is only available for machine charms. When it's run for a K8s charm, Juju logs an error message but doesn't return an response with the error (or have a non-zero exit), which means that we have to fail silently.

Adds a corresponding `reboot()` method to the harness backend.

From the point of view of the harness, we assume that everything is essentially the same after rebooting without `--now`, so this is a no-op. The most likely change that would impact the harness would be a leadership change, but it seems better to let the Charmer model that explicitly in their tests.

When rebooting with `--now`, what should happen is that everything after that point in the event handler is skipped, there's a simulated reboot (again, roughly a no-op), and then the event is re-emitted. I went down a long rabbit-hole trying to get the harness to do this, but it ended up way too messy for minimal benefit. Instead, a `SystemExit` exception is raised and the Charmer is expected to handle it in their tests, for example:

```python
def test_installation(self):
    self.harness.begin()
    with self.assertRaises(SystemExit):
        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.
```

Fixes #929

CHANGES.md

@@ -1,6 +1,9 @@
 # 2.8.0
 
+* 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
 
 # 2.7.0
 

ops/model.py

@@ -26,6 +26,7 @@
 import shutil
 import stat
 import subprocess
+import sys
 import tempfile
 import time
 import typing
@@ -686,6 +687,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.
@@ -709,6 +712,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:
@@ -3400,6 +3429,16 @@ 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")
+            # 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")
+
 
 class _ModelBackendValidator:
     """Provides facilities for validating inputs and formatting them for model backends."""

ops/testing.py

@@ -1741,6 +1741,11 @@ def handle_timeout(args: testing.ExecArgs) -> int:
             handler=(lambda _: result) if handler is None else handler  # type: ignore
         )
 
+    @property
+    def reboot_count(self) -> int:
+        """Number of times the charm has called :meth:`ops.Unit.reboot`."""
+        return self._backend._reboot_count
+
 
 def _get_app_or_unit_name(app_or_unit: AppUnitOrName) -> str:
     """Return name of given application or unit (return strings directly)."""
@@ -1957,6 +1962,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._reboot_count = 0
 
     def _validate_relation_access(self, relation_name: str, relations: List[model.Relation]):
         """Ensures that the named relation exists/has been added.
@@ -2517,6 +2523,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._reboot_count += 1
+        if not now:
+            return
+        # This should exit, reboot, and re-emit the event, but we'll need the caller
+        # to handle everything after the exit.
+        raise SystemExit()
+
 
 @_copy_docstrings(pebble.ExecProcess)
 class _TestingExecProcess:

test/test_model.py

@@ -3596,5 +3596,28 @@ 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', ''],
+        ])
+        with self.assertRaises(SystemExit):
+            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

@@ -3345,6 +3345,37 @@ 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.assertEqual(harness.reboot_count, 0)
+        backend = harness._backend
+        backend.reboot()
+        self.assertEqual(harness.reboot_count, 1)
+        with self.assertRaises(SystemExit):
+            backend.reboot(now=True)
+        self.assertEqual(harness.reboot_count, 2)
+        harness.begin()
+        with self.assertRaises(SystemExit):
+            harness.charm.on.install.emit()
+        self.assertEqual(harness.reboot_count, 3)
+        harness.charm.on.remove.emit()
+        self.assertEqual(harness.reboot_count, 4)
+
 
 class _TestingPebbleClientMixin:
     def get_testing_client(self):