Skip to content

Commit

Permalink
Expose juju-reboot [--now].
Browse files Browse the repository at this point in the history
  • Loading branch information
tonyandrewmeyer committed Oct 10, 2023
1 parent e762ef8 commit 8aec383
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 0 deletions.
32 changes: 32 additions & 0 deletions ops/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down
12 changes: 12 additions & 0 deletions ops/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
22 changes: 22 additions & 0 deletions test/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
26 changes: 26 additions & 0 deletions test/test_testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 8aec383

Please sign in to comment.