Skip to content

Commit

Permalink
refactor!: use ops._main._Manager in Scenario (canonical#1491)
Browse files Browse the repository at this point in the history
The ops._main._Manager class started out as the Ops class in Scenario,
but went through significant modifications before it was added in ops.
This PR merges those changes back into Scenario, so that it's also using
the `_Manager` class. Other than generally simplifying the code, this
also means that Scenario is even more consistent with using ops in
production.

This includes some small refactoring in ops to make this cleaner:
* `ops._main._emit_charm_event` is only used by `_Manager`, so logically
belongs inside that class. This also makes it much simpler to have a
subclass override the behaviour.
* Allow passing a `JujuContext` to `_Manager`, so that it can be
populated by something other than `os.environ`.
* Allow `_Manager` subclasses to override how the charm metadata is
loaded.
* Allow `_Manager` subclasses to execute code after committing but
before the framework is closed.
* Support a dispatch path that uses `_` rather than `-` in event names
(Scenario consistently uses `_`, even though this would always be `-`
from Juju - feel free to push back on this, requiring Scenario to undo
the `_` conversion when populating the context).

The Scenario `runtime` and `ops_main_mock` modules are renamed to
properly indicate that everything inside of them is an implementation
detail and should be considered private by Scenario users. This is a
breaking change but we're considering it a bug fix and not one that
requires a major release, given that the documentation and top-level
namespace were already clear on this. Note that the first commit does
the rename and nothing else - but GitHub shows only one of the files as
a rename (I think because too much content has then changed). It might
be easier to read through the commits separately when reviewing.

Significant Scenario refactoring:
* UnitStateDB is moved from `runtime` to `ops_main_mock`, and it no
longer (re) opens the SQLiteStorage - it instead takes a reference to
the one that's in the framework, so that we're consistently using the
same underlying storage, without relying on the filename being
consistent (this eases moving to `:memory:` in the future).
* As far as I can tell, `runtime._OpsMainContext` is completely unused
and missed being removed in previous refactoring, so it's entirely
deleted.
* Loading and storing deferred events and stored state into the state
database is moved from the runtime to the manager subclass.
* The `Ops` class survives, but is now a subclass of the
`ops._main._Manager` class. A reasonable amount of functionality is now
removed (because the parent handles it), and the remaining code is
specifically for handling the testing framework.

A few type hints are added in some tests where I was working on fixing
the test and it helped to do that.

Fixes canonical#1425
  • Loading branch information
tonyandrewmeyer authored Dec 9, 2024
1 parent ad85c7f commit 0ad731a
Show file tree
Hide file tree
Showing 13 changed files with 294 additions and 483 deletions.
67 changes: 41 additions & 26 deletions ops/_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,28 +112,6 @@ def _setup_event_links(charm_dir: Path, charm: 'ops.charm.CharmBase', juju_conte
_create_event_link(charm, bound_event, link_to)


def _emit_charm_event(charm: 'ops.charm.CharmBase', event_name: str, juju_context: _JujuContext):
"""Emits a charm event based on a Juju event name.
Args:
charm: A charm instance to emit an event from.
event_name: A Juju event name to emit on a charm.
juju_context: An instance of the _JujuContext class.
"""
event_to_emit = None
try:
event_to_emit = getattr(charm.on, event_name)
except AttributeError:
logger.debug('Event %s not defined for %s.', event_name, charm)

# If the event is not supported by the charm implementation, do
# not error out or try to emit it. This is to support rollbacks.
if event_to_emit is not None:
args, kwargs = _get_event_args(charm, event_to_emit, juju_context)
logger.debug('Emitting Juju event %s.', event_name)
event_to_emit.emit(*args, **kwargs)


def _get_event_args(
charm: 'ops.charm.CharmBase',
bound_event: 'ops.framework.BoundEvent',
Expand Down Expand Up @@ -401,8 +379,11 @@ def __init__(
model_backend: Optional[ops.model._ModelBackend] = None,
use_juju_for_storage: Optional[bool] = None,
charm_state_path: str = CHARM_STATE_FILE,
juju_context: Optional[_JujuContext] = None,
):
self._juju_context = _JujuContext.from_dict(os.environ)
if juju_context is None:
juju_context = _JujuContext.from_dict(os.environ)
self._juju_context = juju_context
self._charm_state_path = charm_state_path
self._charm_class = charm_class
if model_backend is None:
Expand All @@ -413,7 +394,7 @@ def __init__(
self._setup_root_logging()

self._charm_root = self._juju_context.charm_dir
self._charm_meta = CharmMeta.from_charm_root(self._charm_root)
self._charm_meta = self._load_charm_meta()
self._use_juju_for_storage = use_juju_for_storage

# Set up dispatcher, framework and charm objects.
Expand All @@ -423,6 +404,9 @@ def __init__(
self.framework = self._make_framework(self.dispatcher)
self.charm = self._make_charm(self.framework, self.dispatcher)

def _load_charm_meta(self):
return CharmMeta.from_charm_root(self._charm_root)

def _make_charm(self, framework: 'ops.framework.Framework', dispatcher: _Dispatcher):
charm = self._charm_class(framework)
dispatcher.ensure_event_links(charm)
Expand Down Expand Up @@ -482,7 +466,7 @@ def _make_framework(self, dispatcher: _Dispatcher):
# If we are in a RelationBroken event, we want to know which relation is
# broken within the model, not only in the event's `.relation` attribute.

if self._juju_context.dispatch_path.endswith('-relation-broken'):
if self._juju_context.dispatch_path.endswith(('-relation-broken', '_relation_broken')):
broken_relation_id = self._juju_context.relation_id
else:
broken_relation_id = None
Expand Down Expand Up @@ -515,19 +499,50 @@ def _emit(self):
self.framework.reemit()

# Emit the Juju event.
_emit_charm_event(self.charm, self.dispatcher.event_name, self._juju_context)
self._emit_charm_event(self.dispatcher.event_name)
# Emit collect-status events.
ops.charm._evaluate_status(self.charm)

def _get_event_to_emit(self, event_name: str) -> Optional[ops.framework.BoundEvent]:
try:
return getattr(self.charm.on, event_name)
except AttributeError:
logger.debug('Event %s not defined for %s.', event_name, self.charm)
return None

def _emit_charm_event(self, event_name: str):
"""Emits a charm event based on a Juju event name.
Args:
charm: A charm instance to emit an event from.
event_name: A Juju event name to emit on a charm.
juju_context: An instance of the _JujuContext class.
"""
event_to_emit = self._get_event_to_emit(event_name)

# If the event is not supported by the charm implementation, do
# not error out or try to emit it. This is to support rollbacks.
if event_to_emit is None:
return

args, kwargs = _get_event_args(self.charm, event_to_emit, self._juju_context)
logger.debug('Emitting Juju event %s.', event_name)
event_to_emit.emit(*args, **kwargs)

def _commit(self):
"""Commit the framework and gracefully teardown."""
self.framework.commit()

def _close(self):
"""Perform any necessary cleanup before the framework is closed."""
# Provided for child classes - nothing needs to be done in the base.

def run(self):
"""Emit and then commit the framework."""
try:
self._emit()
self._commit()
self._close()
finally:
self.framework.close()

Expand Down
2 changes: 1 addition & 1 deletion ops/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,8 +177,8 @@
# monkeypatch it in, so that the ops.testing.ActionFailed object is the
# one that we expect, even if people are mixing Harness and Scenario.
# https://github.com/canonical/ops-scenario/issues/201
import scenario._runtime as _runtime
import scenario.context as _context
import scenario.runtime as _runtime

_context.ActionFailed = ActionFailed # type: ignore[reportPrivateImportUsage]
_runtime.ActionFailed = ActionFailed # type: ignore[reportPrivateImportUsage]
Expand Down
6 changes: 3 additions & 3 deletions test/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def __init__(


@patch('ops._main.setup_root_logging', new=lambda *a, **kw: None) # type: ignore
@patch('ops._main._emit_charm_event', new=lambda *a, **kw: None) # type: ignore
@patch('ops._main._Manager._emit_charm_event', new=lambda *a, **kw: None) # type: ignore
@patch('ops.charm._evaluate_status', new=lambda *a, **kw: None) # type: ignore
class TestCharmInit:
@patch('sys.stderr', new_callable=io.StringIO)
Expand Down Expand Up @@ -235,11 +235,11 @@ def __init__(self, framework: ops.Framework):
dispatch.chmod(0o755)

with patch.dict(os.environ, fake_environ):
with patch('ops._main._emit_charm_event') as mock_charm_event:
with patch('ops._main._Manager._emit_charm_event') as mock_charm_event:
ops.main(MyCharm)

assert mock_charm_event.call_count == 1
return mock_charm_event.call_args[0][1]
return mock_charm_event.call_args[0][0]

def test_most_legacy(self):
"""Without dispatch, sys.argv[0] is used."""
Expand Down
2 changes: 1 addition & 1 deletion test/test_main_invocation.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
@pytest.fixture
def charm_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
monkeypatch.setattr('sys.argv', ('hooks/install',))
monkeypatch.setattr('ops._main._emit_charm_event', Mock())
monkeypatch.setattr('ops._main._Manager._emit_charm_event', Mock())
monkeypatch.setattr('ops._main._Manager._setup_root_logging', Mock())
monkeypatch.setattr('ops.charm._evaluate_status', Mock())
monkeypatch.setenv('JUJU_CHARM_DIR', str(tmp_path))
Expand Down
7 changes: 6 additions & 1 deletion testing/src/scenario/_consistency_checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
)

from .errors import InconsistentScenarioError
from .runtime import logger as scenario_logger
from ._runtime import logger as scenario_logger
from .state import (
CharmType,
PeerRelation,
Expand Down Expand Up @@ -179,6 +179,11 @@ def check_event_consistency(
# skip everything here. Perhaps in the future, custom events could
# optionally include some sort of state metadata that made testing
# consistency possible?
warnings.append(
"this is a custom event; if its name makes it look like a builtin one "
"(for example, a relation event, or a workload event), you might get some false-negative "
"consistency checks.",
)
return Results(errors, warnings)

if event._is_relation_event:
Expand Down
197 changes: 197 additions & 0 deletions testing/src/scenario/_ops_main_mock.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,197 @@
#!/usr/bin/env python3
# Copyright 2023 Canonical Ltd.
# See LICENSE file for licensing details.

import dataclasses
import marshal
import re
import sys
from typing import TYPE_CHECKING, Any, Dict, FrozenSet, List, Sequence, Set

import ops
import ops.jujucontext
import ops.storage

from ops.framework import _event_regex
from ops._main import _Dispatcher, _Manager
from ops._main import logger as ops_logger

from .errors import BadOwnerPath, NoObserverError
from .logger import logger as scenario_logger
from .mocking import _MockModelBackend
from .state import CharmType, StoredState, DeferredEvent

if TYPE_CHECKING: # pragma: no cover
from .context import Context
from .state import State, _CharmSpec, _Event

EVENT_REGEX = re.compile(_event_regex)
STORED_STATE_REGEX = re.compile(
r"((?P<owner_path>.*)\/)?(?P<_data_type_name>\D+)\[(?P<name>.*)\]",
)

logger = scenario_logger.getChild("ops_main_mock")

# pyright: reportPrivateUsage=false


class UnitStateDB:
"""Wraps the unit-state database with convenience methods for adjusting the state."""

def __init__(self, underlying_store: ops.storage.SQLiteStorage):
self._db = underlying_store

def get_stored_states(self) -> FrozenSet["StoredState"]:
"""Load any StoredState data structures from the db."""
db = self._db
stored_states: Set[StoredState] = set()
for handle_path in db.list_snapshots():
if not EVENT_REGEX.match(handle_path) and (
match := STORED_STATE_REGEX.match(handle_path)
):
stored_state_snapshot = db.load_snapshot(handle_path)
kwargs = match.groupdict()
sst = StoredState(content=stored_state_snapshot, **kwargs)
stored_states.add(sst)

return frozenset(stored_states)

def get_deferred_events(self) -> List["DeferredEvent"]:
"""Load any DeferredEvent data structures from the db."""
db = self._db
deferred: List[DeferredEvent] = []
for handle_path in db.list_snapshots():
if EVENT_REGEX.match(handle_path):
notices = db.notices(handle_path)
for handle, owner, observer in notices:
try:
snapshot_data = db.load_snapshot(handle)
except ops.storage.NoSnapshotError:
snapshot_data: Dict[str, Any] = {}

event = DeferredEvent(
handle_path=handle,
owner=owner,
observer=observer,
snapshot_data=snapshot_data,
)
deferred.append(event)

return deferred

def apply_state(self, state: "State"):
"""Add DeferredEvent and StoredState from this State instance to the storage."""
db = self._db
for event in state.deferred:
db.save_notice(event.handle_path, event.owner, event.observer)
try:
marshal.dumps(event.snapshot_data)
except ValueError as e:
raise ValueError(
f"unable to save the data for {event}, it must contain only simple types.",
) from e
db.save_snapshot(event.handle_path, event.snapshot_data)

for stored_state in state.stored_states:
db.save_snapshot(stored_state._handle_path, stored_state.content)


class Ops(_Manager):
"""Class to manage stepping through ops setup, event emission and framework commit."""

def __init__(
self,
state: "State",
event: "_Event",
context: "Context[CharmType]",
charm_spec: "_CharmSpec[CharmType]",
juju_context: ops.jujucontext._JujuContext,
):
self.state = state
self.event = event
self.context = context
self.charm_spec = charm_spec
self.store = None

model_backend = _MockModelBackend(
state=state,
event=event,
context=context,
charm_spec=charm_spec,
juju_context=juju_context,
)

super().__init__(
self.charm_spec.charm_type, model_backend, juju_context=juju_context
)

def _load_charm_meta(self):
metadata = (self._charm_root / "metadata.yaml").read_text()
actions_meta = self._charm_root / "actions.yaml"
if actions_meta.exists():
actions_metadata = actions_meta.read_text()
else:
actions_metadata = None

return ops.CharmMeta.from_yaml(metadata, actions_metadata)

def _setup_root_logging(self):
# Ops sets sys.excepthook to go to Juju's debug-log, but that's not
# useful in a testing context, so we reset it here.
super()._setup_root_logging()
sys.excepthook = sys.__excepthook__

def _make_storage(self, _: _Dispatcher):
# TODO: add use_juju_for_storage support
# TODO: Pass a charm_state_path that is ':memory:' when appropriate.
charm_state_path = self._charm_root / self._charm_state_path
storage = ops.storage.SQLiteStorage(charm_state_path)
logger.info("Copying input state to storage.")
self.store = UnitStateDB(storage)
self.store.apply_state(self.state)
return storage

def _get_event_to_emit(self, event_name: str):
owner = (
self._get_owner(self.charm, self.event.owner_path)
if self.event
else self.charm.on
)

try:
event_to_emit = getattr(owner, event_name)
except AttributeError:
ops_logger.debug("Event %s not defined for %s.", event_name, self.charm)
raise NoObserverError(
f"Cannot fire {event_name!r} on {owner}: "
f"invalid event (not on charm.on).",
)
return event_to_emit

@staticmethod
def _get_owner(root: Any, path: Sequence[str]) -> ops.ObjectEvents:
"""Walk path on root to an ObjectEvents instance."""
obj = root
for step in path:
try:
obj = getattr(obj, step)
except AttributeError:
raise BadOwnerPath(
f"event_owner_path {path!r} invalid: {step!r} leads to nowhere.",
)
if not isinstance(obj, ops.ObjectEvents):
raise BadOwnerPath(
f"event_owner_path {path!r} invalid: does not lead to "
f"an ObjectEvents instance.",
)
return obj

def _close(self):
"""Now that we're done processing this event, read the charm state and expose it."""
logger.info("Copying storage to output state.")
assert self.store is not None
deferred = self.store.get_deferred_events()
stored_state = self.store.get_stored_states()
self.state = dataclasses.replace(
self.state, deferred=deferred, stored_states=stored_state
)
Loading

0 comments on commit 0ad731a

Please sign in to comment.