diff --git a/ops/model.py b/ops/model.py index 1e8206ab8..3aafc8d41 100644 --- a/ops/model.py +++ b/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: @@ -3349,6 +3375,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.""" diff --git a/ops/testing.py b/ops/testing.py index 57a0c52a0..128b52d61 100755 --- a/ops/testing.py +++ b/ops/testing.py @@ -1882,6 +1882,10 @@ class _Secret: grants: Dict[int, Set[str]] = dataclasses.field(default_factory=dict) +class RebootingMachineError(Exception): + """Raised when the machine would reboot.""" + + @_copy_docstrings(model._ModelBackend) @_record_calls class _TestingModelBackend: @@ -2503,6 +2507,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: diff --git a/test/test_model.py b/test/test_model.py index 365ddff2a..76d001f47 100755 --- a/test/test_model.py +++ b/test/test_model.py @@ -3596,5 +3596,27 @@ def test_set_ports_noop(self): ]) +class TestContainer(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() diff --git a/test/test_testing.py b/test/test_testing.py index 9d4f9225e..57cb90ede 100644 --- a/test/test_testing.py +++ b/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):