Repository: canonical/operator
Pull request #1085 (2023-12-01 09:12:35)

chore: Refactor main.py to add a new _Manager class
With this I hope to contribute some of the abstractions I developed for Scenario's `ops.main` mock back to the framework.

It's **not a functional change**, that is, ops should be functioning in exactly the same way.

It is however IMHO a positive change for the following reasons:
- tidier, more object-oriented. `main.py` was hard to understand and extend
- (therefore) easier to play with in e.g. testing frameworks or experimental modules wishing to modify or extend its functionality
- easier to maintain

The core contributions are:
- a `_Manager` class encapsulating the management of the ops lifecycle, from instantiation of Framework and charm instances to storage creation, event emission, and storage commit.
- adding the `CharmMeta.from_charm_root` method to load a `CharmMeta` instance from the relevant `*.yaml` files

PietroPasotti: (test failures seem to be related to the breaking changes recently introduced in harness' secret model
PietroPasotti: also charmcraft pack has an unrelated bug https://matrix.to/#/!LhFxJIPEcCacgdMghH:ubuntu.com/$ah10i3Ckd9YsahSr9xROqQNOTfVK_VrTYDK3oszGQGE?via=ubuntu.com&via=matrix.org hence the integration failures
tonyandrewmeyer: > (TBD if we can use dataclasses or we need to support older python versions, if so we could fall back to a namedtuple) that encapsulates the 'charm module' including its metadata yamls and repo root.

For ops itself, we explicitly support 3.8+, so dataclasses are fine (and there are already a couple of uses in ops). If the concern is about experiments (etc) using this class then I would say we do not want to support anything older than ops itself supports.
benhoyt: Thanks @PietroPasotti. Tony and I will take a look over this and discuss tomorrow and then get back to you with some thoughts or next steps.
PietroPasotti: will get back to this next week when I have more time. Thanks for the feedback.
PietroPasotti: @tonyandrewmeyer @benhoyt sorry for taking so long to get back on this, I was on holiday in december and forgot it existed. Addressed all comments and rebased.
benhoyt: @PietroPasotti Looks like `test_crash_action` is failing:

```
FAILED test/test_main.py::TestMainWithDispatch::test_crash_action - AssertionError: 'KeyError' not found in ''
```

PietroPasotti: Ok I'm starting to see the light on _CharmSpec kind of being the same thing as CharmMeta. I somewhat thought charmmeta was supposed to wrap only the metadata (as in what's in `metadata.yaml`, but I guess that's not the best way to look at it.
tonyandrewmeyer: I went ahead and made the small tweaks - no point bouncing this around again for that. I'll just run this against the charm tests as a sanity check that it's not going to break anything, and then we'll be good to go.

Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 5faa209..3c3716d 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -2,6 +2,7 @@
 
 * Updated Pebble Notices `get_notices` parameter name to `users=all` (previously `select=all`).
 * Added `Model.get_cloud_spec` which uses the `credential-get` hook tool to get details of the cloud where the model is deployed.
+* Refactored main.py, creating a new `_Manager` class.
 
 # 2.11.0
 
diff --git a/ops/charm.py b/ops/charm.py
index eae71a2..fb687b4 100644
--- a/ops/charm.py
+++ b/ops/charm.py
@@ -1373,6 +1373,24 @@ class CharmMeta:
         self.containers = {name: ContainerMeta(name, container)
                            for name, container in raw_.get('containers', {}).items()}
 
+    @staticmethod
+    def from_charm_root(charm_root: Union[pathlib.Path, str]):
+        """Initialise CharmMeta from the path to a charm repository root folder."""
+        _charm_root = pathlib.Path(charm_root)
+        metadata_path = _charm_root / "metadata.yaml"
+
+        with metadata_path.open() as f:
+            meta = yaml.safe_load(f.read())
+
+        actions = None
+
+        actions_path = _charm_root / "actions.yaml"
+        if actions_path.exists():
+            with actions_path.open() as f:
+                actions = yaml.safe_load(f.read())
+
+        return CharmMeta(meta, actions)
+
     def _load_links(self, raw: Dict[str, Any]):
         websites = raw.get('website', [])
         if not websites and 'links' in raw:
diff --git a/ops/main.py b/ops/main.py
index b3c013b..7c60333 100644
--- a/ops/main.py
+++ b/ops/main.py
@@ -18,7 +18,6 @@ Note that this module is callable, and calls the :func:`ops.main.main` function.
 This is so that :code:`import ops` followed by :code:`ops.main(MyCharm)` works
 as expected.
 """
-
 import logging
 import os
 import shutil
@@ -294,10 +293,10 @@ class _Dispatcher:
             subprocess.run(argv, check=True)
         except subprocess.CalledProcessError as e:
             logger.warning("Legacy %s exited with status %d.", self._dispatch_path, e.returncode)
-            sys.exit(e.returncode)
+            raise _Abort(e.returncode) from e
         except OSError as e:
             logger.warning("Unable to run legacy %s: %s", self._dispatch_path, e)
-            sys.exit(1)
+            raise _Abort(1) from e
         else:
             logger.debug("Legacy %s exited with status 0.", self._dispatch_path)
 
@@ -330,7 +329,7 @@ class _Dispatcher:
 
         if 'OPERATOR_DISPATCH' in os.environ:
             logger.debug("Charm called itself via %s.", self._dispatch_path)
-            sys.exit(0)
+            raise _Abort(0)
         os.environ['OPERATOR_DISPATCH'] = '1'
 
         self.is_dispatch_aware = True
@@ -369,97 +368,182 @@ def _should_use_controller_storage(db_path: Path, meta: CharmMeta) -> bool:
         return False
 
 
-def main(charm_class: Type[ops.charm.CharmBase],
-         use_juju_for_storage: Optional[bool] = None):
-    """Setup the charm and dispatch the observed event.
+class _Abort(Exception):  # noqa: N818
+    """Raised when something happens that should interrupt ops execution."""
 
-    The event name is based on the way this executable was called (argv[0]).
+    def __init__(self, exit_code: int):
+        super().__init__()
+        self.exit_code = exit_code
 
-    Args:
-        charm_class: the charm class to instantiate and receive the event.
-        use_juju_for_storage: whether to use controller-side storage. If not specified
-            then kubernetes charms that haven't previously used local storage and that
-            are running on a new enough Juju default to controller-side storage,
-            otherwise local storage is used.
+
+class _Manager:
+    """Initialises the Framework and manages the lifecycle of a charm.
+
+    Running _Manager consists of three main steps:
+    - setup: initialise the following from JUJU_* environment variables:
+      - the Framework (hook tool wrappers)
+      - the storage backend
+      - the event that Juju is emitting on us
+      - the charm instance (user-facing)
+    - emit: core user-facing lifecycle step. Consists of:
+      - re-emit any deferred events found in the storage
+      - emit the Juju event to the charm
+        - emit any custom events emitted by the charm during this phase
+      - emit  the ``collect-status`` events
+    - commit: responsible for:
+      - store any events deferred throughout this execution
+      - graceful teardown of the storage
     """
-    charm_dir = _get_charm_dir()
-
-    model_backend = ops.model._ModelBackend()
-    debug = ('JUJU_DEBUG' in os.environ)
-    # For actions, there is a communication channel with the user running the
-    # action, so we want to send exception details through stderr, rather than
-    # only to juju-log as normal.
-    handling_action = ('JUJU_ACTION_NAME' in os.environ)
-    setup_root_logging(model_backend, debug=debug, exc_stderr=handling_action)
-    logger.debug("ops %s up and running.", ops.__version__)  # type:ignore
-
-    dispatcher = _Dispatcher(charm_dir)
-    dispatcher.run_any_legacy_hook()
-
-    metadata = (charm_dir / 'metadata.yaml').read_text()
-    actions_meta = charm_dir / 'actions.yaml'
-    actions_metadata = actions_meta.read_text() if actions_meta.exists() else None
-
-    # 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 os.environ.get('JUJU_DISPATCH_PATH', '').endswith('-relation-broken'):
-        broken_relation_id = _get_juju_relation_id()
-    else:
-        broken_relation_id = None
-    meta = CharmMeta.from_yaml(metadata, actions_metadata)
-    model = ops.model.Model(meta, model_backend, broken_relation_id=broken_relation_id)
-
-    charm_state_path = charm_dir / CHARM_STATE_FILE
-
-    if use_juju_for_storage and not ops.storage.juju_backend_available():
-        # raise an exception; the charm is broken and needs fixing.
-        msg = 'charm set use_juju_for_storage=True, but Juju version {} does not support it'
-        raise RuntimeError(msg.format(JujuVersion.from_environ()))
-
-    if use_juju_for_storage is None:
-        use_juju_for_storage = _should_use_controller_storage(charm_state_path, meta)
-    elif use_juju_for_storage:
-        warnings.warn("Controller storage is deprecated; it's intended for "
-                      "podspec charms and will be removed in a future release.",
-                      category=DeprecationWarning)
-
-    if use_juju_for_storage:
-        if dispatcher.is_restricted_context():
+
+    def __init__(
+            self,
+            charm_class: Type["ops.charm.CharmBase"],
+            model_backend: Optional[ops.model._ModelBackend] = None,
+            use_juju_for_storage: Optional[bool] = None,
+            charm_state_path: str = CHARM_STATE_FILE
+    ):
+
+        self._charm_state_path = charm_state_path
+        self._charm_class = charm_class
+        if model_backend is None:
+            model_backend = ops.model._ModelBackend()
+        self._model_backend = model_backend
+
+        # Do this as early as possible to be sure to catch the most logs.
+        self._setup_root_logging()
+
+        self._charm_root = _get_charm_dir()
+        self._charm_meta = CharmMeta.from_charm_root(self._charm_root)
+        self._use_juju_for_storage = use_juju_for_storage
+
+        # Set up dispatcher, framework and charm objects.
+        self.dispatcher = _Dispatcher(self._charm_root)
+        self.dispatcher.run_any_legacy_hook()
+
+        self.framework = self._make_framework(self.dispatcher)
+        self.charm = self._make_charm(self.framework, self.dispatcher)
+
+    def _make_charm(self, framework: "ops.framework.Framework", dispatcher: _Dispatcher):
+        charm = self._charm_class(framework)
+        dispatcher.ensure_event_links(charm)
+        return charm
+
+    def _setup_root_logging(self):
+        debug = "JUJU_DEBUG" in os.environ
+        # For actions, there is a communication channel with the user running the
+        # action, so we want to send exception details through stderr, rather than
+        # only to juju-log as normal.
+        handling_action = 'JUJU_ACTION_NAME' in os.environ
+        setup_root_logging(self._model_backend, debug=debug, exc_stderr=handling_action)
+
+        logger.debug("ops %s up and running.", ops.__version__)  # type:ignore
+
+    def _make_storage(self, dispatcher: _Dispatcher):
+        charm_state_path = self._charm_root / self._charm_state_path
+
+        use_juju_for_storage = self._use_juju_for_storage
+        if use_juju_for_storage and not ops.storage.juju_backend_available():
+            # raise an exception; the charm is broken and needs fixing.
+            msg = 'charm set use_juju_for_storage=True, but Juju version {} does not support it'
+            raise RuntimeError(msg.format(JujuVersion.from_environ()))
+
+        if use_juju_for_storage is None:
+            use_juju_for_storage = _should_use_controller_storage(
+                charm_state_path,
+                self._charm_meta
+            )
+        elif use_juju_for_storage:
+            warnings.warn("Controller storage is deprecated; it's intended for "
+                          "podspec charms and will be removed in a future release.",
+                          category=DeprecationWarning)
+
+        if use_juju_for_storage and dispatcher.is_restricted_context():
             # TODO: jam 2020-06-30 This unconditionally avoids running a collect metrics event
             #  Though we eventually expect that Juju will run collect-metrics in a
-            #  non-restricted context. Once we can determine that we are running collect-metrics
-            #  in a non-restricted context, we should fire the event as normal.
+            #  non-restricted context. Once we can determine that we are running
+            #  collect-metrics in a non-restricted context, we should fire the event as normal.
             logger.debug('"%s" is not supported when using Juju for storage\n'
                          'see: https://github.com/canonical/operator/issues/348',
                          dispatcher.event_name)
             # Note that we don't exit nonzero, because that would cause Juju to rerun the hook
-            return
-        store = ops.storage.JujuStorage()
-    else:
-        store = ops.storage.SQLiteStorage(charm_state_path)
-    framework = ops.framework.Framework(store, charm_dir, meta, model,
-                                        event_name=dispatcher.event_name)
-    framework.set_breakpointhook()
-    try:
-        charm = charm_class(framework)
-        dispatcher.ensure_event_links(charm)
+            raise _Abort(0)
 
+        if self._use_juju_for_storage:
+            store = ops.storage.JujuStorage()
+        else:
+            store = ops.storage.SQLiteStorage(charm_state_path)
+        return store
+
+    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 os.environ.get('JUJU_DISPATCH_PATH', '').endswith('-relation-broken'):
+            broken_relation_id = _get_juju_relation_id()
+        else:
+            broken_relation_id = None
+
+        model = ops.model.Model(self._charm_meta, self._model_backend,
+                                broken_relation_id=broken_relation_id)
+        store = self._make_storage(dispatcher)
+        framework = ops.framework.Framework(store, self._charm_root, self._charm_meta, model,
+                                            event_name=dispatcher.event_name)
+        framework.set_breakpointhook()
+        return framework
+
+    def _emit(self):
+        """Emit the event on the charm."""
         # TODO: Remove the collect_metrics check below as soon as the relevant
         #       Juju changes are made. Also adjust the docstring on
         #       EventBase.defer().
         #
         # Skip reemission of deferred events for collect-metrics events because
         # they do not have the full access to all hook tools.
-        if not dispatcher.is_restricted_context():
-            framework.reemit()
+        if not self.dispatcher.is_restricted_context():
+            # Re-emit any deferred events from the previous run.
+            self.framework.reemit()
+
+        # Emit the Juju event.
+        _emit_charm_event(self.charm, self.dispatcher.event_name)
+        # Emit collect-status events.
+        ops.charm._evaluate_status(self.charm)
 
-        _emit_charm_event(charm, dispatcher.event_name)
+    def _commit(self):
+        """Commit the framework and gracefully teardown."""
+        self.framework.commit()
 
-        ops.charm._evaluate_status(charm)
+    def run(self):
+        """Emit and then commit the framework."""
+        try:
+            self._emit()
+            self._commit()
+        finally:
+            self.framework.close()
+
+
+def main(charm_class: Type[ops.charm.CharmBase],
+         use_juju_for_storage: Optional[bool] = None):
+    """Set up the charm and dispatch the observed event.
+
+    The event name is based on the way this executable was called (argv[0]).
+
+    Args:
+        charm_class: the charm class to instantiate and receive the event.
+        use_juju_for_storage: whether to use controller-side storage. If not specified
+            then Kubernetes charms that haven't previously used local storage and that
+            are running on a new enough Juju default to controller-side storage,
+            otherwise local storage is used.
+    """
+    try:
+        manager = _Manager(
+            charm_class,
+            use_juju_for_storage=use_juju_for_storage)
 
-        framework.commit()
-    finally:
-        framework.close()
+        manager.run()
+    except _Abort as e:
+        sys.exit(e.exit_code)
 
 
 # Make this module callable and call main(), so that "import ops" and then


Test patch:
diff --git a/test/test_charm.py b/test/test_charm.py
index 9767588..b471b37 100644
--- a/test/test_charm.py
+++ b/test/test_charm.py
@@ -13,12 +13,15 @@
 # limitations under the License.
 import functools
 import os
+import pathlib
 import shutil
 import tempfile
 import typing
 import unittest
 from pathlib import Path
 
+import yaml
+
 import ops
 import ops.charm
 from ops.model import ModelError, _ModelBackend
@@ -473,6 +476,45 @@ start:
   additionalProperties: false
 ''')
 
+    def test_meta_from_charm_root(self):
+        with tempfile.TemporaryDirectory() as d:
+            td = pathlib.Path(d)
+            (td / 'metadata.yaml').write_text(
+                yaml.safe_dump(
+                    {"name": "bob",
+                        "requires": {
+                            "foo":
+                                {"interface": "bar"}
+                        }}))
+            meta = ops.CharmMeta.from_charm_root(td)
+            self.assertEqual(meta.name, "bob")
+            self.assertEqual(meta.requires['foo'].interface_name, "bar")
+
+    def test_actions_from_charm_root(self):
+        with tempfile.TemporaryDirectory() as d:
+            td = pathlib.Path(d)
+            (td / 'actions.yaml').write_text(
+                yaml.safe_dump(
+                    {"foo": {
+                        "description": "foos the bar",
+                        "additionalProperties": False
+                    }}
+                )
+            )
+            (td / 'metadata.yaml').write_text(
+                yaml.safe_dump(
+                    {"name": "bob",
+                        "requires": {
+                            "foo":
+                                {"interface": "bar"}
+                        }}))
+
+            meta = ops.CharmMeta.from_charm_root(td)
+            self.assertEqual(meta.name, "bob")
+            self.assertEqual(meta.requires['foo'].interface_name, "bar")
+            self.assertFalse(meta.actions['foo'].additional_properties)
+            self.assertEqual(meta.actions['foo'].description, "foos the bar")
+
     def _setup_test_action(self):
         fake_script(self, 'action-get', """echo '{"foo-name": "name", "silent": true}'""")
         fake_script(self, 'action-set', "")
diff --git a/test/test_main.py b/test/test_main.py
index 83f541b..c22d0b5 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -29,7 +29,7 @@ from pathlib import Path
 from unittest.mock import patch
 
 import ops
-from ops.main import CHARM_STATE_FILE, _should_use_controller_storage
+from ops.main import _should_use_controller_storage
 from ops.storage import SQLiteStorage
 
 from .charms.test_main.src.charm import MyCharmEvents
@@ -196,7 +196,7 @@ class CharmInitTestCase(unittest.TestCase):
 
 
 @patch('sys.argv', new=("hooks/config-changed",))
-@patch('ops.main.setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
+@patch('ops.main._Manager._setup_root_logging', new=lambda *a, **kw: None)  # type: ignore
 @patch('ops.charm._evaluate_status', new=lambda *a, **kw: None)  # type: ignore
 class TestDispatch(unittest.TestCase):
     def _check(self, *, with_dispatch: bool = False, dispatch_path: str = ''):
@@ -316,7 +316,7 @@ class _TestMain(abc.ABC):
         self._tmpdir = Path(tempfile.mkdtemp(prefix='tmp-ops-test-')).resolve()
         self.addCleanup(shutil.rmtree, str(self._tmpdir))
         self.JUJU_CHARM_DIR = self._tmpdir / 'test_main'
-        self.CHARM_STATE_FILE = self.JUJU_CHARM_DIR / CHARM_STATE_FILE
+        self._charm_state_file = self.JUJU_CHARM_DIR / '.unit-state.db'
         self.hooks_dir = self.JUJU_CHARM_DIR / 'hooks'
         charm_path = str(self.JUJU_CHARM_DIR / 'src/charm.py')
         self.charm_exec_path = os.path.relpath(charm_path, str(self.hooks_dir))
@@ -348,8 +348,8 @@ class _TestMain(abc.ABC):
     def _read_and_clear_state(self,
                               event_name: str) -> typing.Union[ops.BoundStoredState,
                                                                ops.StoredStateData]:
-        if self.CHARM_STATE_FILE.stat().st_size:
-            storage = SQLiteStorage(self.CHARM_STATE_FILE)
+        if self._charm_state_file.stat().st_size:
+            storage = SQLiteStorage(self._charm_state_file)
             with (self.JUJU_CHARM_DIR / 'metadata.yaml').open() as m:
                 af = (self.JUJU_CHARM_DIR / 'actions.yaml')
                 if af.exists():