Harness Extensions

'''This is a library providing a utility for unittesting code that's meant to be used in charms.

So not charm code per se, but e.g. library code, extensions, etc...

Basic usage:

>>> from charms.harness_extensions.v0.charm_tester import harness_factory
>>> def test_lib():
>>>     harness = harness_factory()
>>>     harness.begin()
>>>     charm = harness.charm
>>>
>>>     @charm.run
>>>     def _initialize(self):
>>>         self.lib = CharmLib(self)
>>>
>>>     @charm.listener(charm.on.my_lib_event)
>>>     def _on_my_lib_event(self, event):
>>>         assert event.foo == 'bar'
>>>
>>>     assert not _on_my_lib_event.called
>>>     harness.do_things_to_trigger_lib_event()
>>>     assert isinstance(_on_my_lib_event.called, MyExpectedEventType)
'''

# The unique Charmhub library identifier, never change it
LIBID = "9634087a8856453db6308a63f61ca8df"

# Increment this major API version when introducing breaking changes
LIBAPI = 0

# Increment this PATCH version before using `charmcraft publish-lib` or reset
# to 0 if you are raising the major API version
LIBPATCH = 2

import typing
from functools import partial

from ops.charm import CharmBase
from ops.framework import BoundEvent, EventBase
from ops.testing import Harness

OptionalYAML = typing.Optional[typing.Union[str, typing.TextIO]]


class _TestCharmABC(CharmBase):
    def get_calls(self, clear: bool = False) -> typing.List[typing.Any]: ...

    def run(self, fn: typing.Callable[[CharmBase], typing.Any]) -> None: ...

    def listener(self, event: str) -> typing.Callable[
        [typing.Callable[[CharmBase], EventBase]], None]: ...

    def register_listener(self, event: BoundEvent,
                          callback: typing.Callable[[CharmBase, BoundEvent], None]): ...


def charm_type_factory() -> typing.Type[CharmBase]:
    class InvokeEvent(EventBase):
        pass

    class TestCharm(CharmBase):
        def __init__(self, framework, key=None):
            super().__init__(framework, key)
            self._callback = None
            self.on.define_event('invoke', InvokeEvent)
            self.framework.observe(self.on.invoke, self._on_invoke)

            self._listeners = {}
            self._listener_calls = []

        def get_calls(self, clear=False):
            calls = self._listener_calls
            if clear:
                self._listener_calls = []
            return calls

        def run(self, fn: typing.Callable[[CharmBase], typing.Any]):
            if self._callback:
                raise RuntimeError('already in a run scope')

            self._callback = partial(fn, self)
            self._invoke()
            self._callback = None

        def _invoke(self, *args):
            self.on.invoke.emit(*args)

        def _on_invoke(self, event):
            self._callback()

        def listener(self, event: str):
            def wrapper(callback):
                self.register_listener(event, callback)
                callback.called = False
                return callback

            return wrapper

        def register_listener(self, event: BoundEvent, callback):
            self._listeners[event.event_kind] = callback
            self.framework.observe(event, self._call_listener)

        def _call_listener(self, evt: EventBase):
            listener = self._listeners[evt.handle.kind]
            self._listener_calls.append(listener)
            listener.called = evt
            listener(self, evt)

    return TestCharm


def harness_factory(meta: OptionalYAML = None,
                    actions: OptionalYAML = None,
                    config: OptionalYAML = None) -> Harness[_TestCharmABC]:
    return Harness(charm_type_factory(), meta=meta, actions=actions, config=config)