storage and entitystatus subclass checks

README.md

@@ -500,7 +500,7 @@ remote_unit_2_is_joining_event = relation.joined_event(remote_unit_id=2)
 remote_unit_2_is_joining_event = Event('foo-relation-changed', relation=relation, relation_remote_unit_id=2)
 ```
 
-## Containers
+# Containers
 
 When testing a kubernetes charm, you can mock container interactions. When using the null state (`State()`), there will
 be no containers. So if the charm were to `self.unit.containers`, it would get back an empty dict.
@@ -586,7 +586,7 @@ need to associate the container with the event is that the Framework uses an env
 pebble-ready event is about (it does not use the event name). Scenario needs that information, similarly, for injecting
 that envvar into the charm's runtime.
 
-### Container filesystem post-mortem
+## Container filesystem post-mortem
 If the charm writes files to a container (to a location you didn't Mount as a temporary folder you have access to), you will be able to inspect them using the `get_filesystem` api.
 
 ```python
@@ -623,7 +623,7 @@ def test_pebble_push():
     assert cfg_file.read_text() == "TEST"
 ```
 
-### `Container.exec` mocks
+## `Container.exec` mocks
 
 `container.exec` is a tad more complicated, but if you get to this low a level of simulation, you probably will have far
 worse issues to deal with. You need to specify, for each possible command the charm might run on the container, what the
@@ -671,6 +671,77 @@ def test_pebble_exec():
     )
 ```
 
+# Storage
+
+If your charm defines `storage` in its metadata, you can use `scenario.state.Storage` to instruct Scenario to make (mocked) filesystem storage available to the charm at runtime.
+
+Using the same `get_filesystem` API as `Container`, you can access the tempdir used by Scenario to mock the filesystem root before and after the scenario runs.
+
+```python
+from scenario import Storage, Context, State
+# some charm with a 'foo' filesystem-type storage defined in metadata.yaml 
+ctx = Context(MyCharm)
+storage = Storage("foo")
+# setup storage with some content
+(storage.get_filesystem(ctx) / "myfile.txt").write_text("helloworld")
+
+with ctx.manager("update-status", State(storage=[storage])) as mgr:
+    foo = mgr.charm.model.storages["foo"][0]
+    loc = foo.location
+    path = loc / "myfile.txt"
+    assert path.exists()
+    assert path.read_text() == "helloworld"
+
+    myfile = loc / "path.py"
+    myfile.write_text("helloworlds")
+
+# post-mortem: inspect fs contents.
+assert (
+    storage.get_filesystem(ctx) / "path.py"
+).read_text() == "helloworlds"
+```
+
+Note that State only wants to know about **attached** storages. A storage which is not attached to the charm can simply be omitted from State and the charm will be none the wiser.
+
+## Storage-add
+
+If a charm requests adding more storage instances while handling some event, you can inspect that from the `Context.requested_storage` API.
+
+```python
+# in MyCharm._on_foo:
+# the charm requests two new "foo" storage instances to be provisioned 
+self.model.storages.request("foo", 2)
+```
+
+From test code, you can inspect that:
+
+```python
+from scenario import Context, State
+
+ctx = Context(MyCharm)
+ctx.run('some-event-that-will-cause_on_foo-to-be-called', State())
+
+# the charm has requested two 'foo' storages to be provisioned
+assert ctx.requested_storages['foo'] == 2
+```
+
+Requesting storages has no other consequence in Scenario. In real life, this request will trigger Juju to provision the storage and execute the charm again with `foo-storage-attached`.
+So a natural follow-up Scenario test suite for this case would be:
+
+```python
+from scenario import Context, State, Storage
+
+ctx = Context(MyCharm)
+foo_0 = Storage('foo')
+# the charm is notified that one of the storages it has requested is ready
+ctx.run(foo_0.attached_event, State(storage=[foo_0]))
+
+foo_1 = Storage('foo')
+# the charm is notified that the other storage is also ready
+ctx.run(foo_1.attached_event, State(storage=[foo_0, foo_1]))
+```
+
+
 # Ports
 
 Since `ops 2.6.0`, charms can invoke the `open-port`, `close-port`, and `opened-ports` hook tools to manage the ports opened on the host vm/container. Using the `State.opened_ports` api, you can: 

scenario/__init__.py

@@ -21,6 +21,7 @@
     Secret,
     State,
     StateValidationError,
+    Storage,
     StoredState,
     SubordinateRelation,
     deferred,
@@ -45,6 +46,7 @@
     "BindAddress",
     "Network",
     "Port",
+    "Storage",
     "StoredState",
     "State",
     "DeferredEvent",

scenario/context.py

@@ -246,6 +246,7 @@ def __init__(
         self.unit_status_history: List["_EntityStatus"] = []
         self.workload_version_history: List[str] = []
         self.emitted_events: List[EventBase] = []
+        self.requested_storages: Dict[str, int] = {}
 
         # set by Runtime.exec() in self._run()
         self._output_state: Optional["State"] = None
@@ -263,13 +264,21 @@ def _get_container_root(self, container_name: str):
         """Get the path to a tempdir where this container's simulated root will live."""
         return Path(self._tmp.name) / "containers" / container_name
 
+    def _get_storage_root(self, name: str, index: int) -> Path:
+        """Get the path to a tempdir where this storage's simulated root will live."""
+        storage_root = Path(self._tmp.name) / "storages" / f"{name}-{index}"
+        # in the case of _get_container_root, _MockPebbleClient will ensure the dir exists.
+        storage_root.mkdir(parents=True, exist_ok=True)
+        return storage_root
+
     def clear(self):
         """Cleanup side effects histories."""
         self.juju_log = []
         self.app_status_history = []
         self.unit_status_history = []
         self.workload_version_history = []
         self.emitted_events = []
+        self.requested_storages = {}
         self._action_logs = []
         self._action_results = None
         self._action_failure = ""

scenario/mocking.py

@@ -6,10 +6,11 @@
 import shutil
 from io import StringIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
 
 from ops import pebble
 from ops.model import (
+    ModelError,
     SecretInfo,
     SecretRotate,
     _format_action_result_dict,
@@ -19,7 +20,7 @@
 from ops.testing import _TestingPebbleClient
 
 from scenario.logger import logger as scenario_logger
-from scenario.state import JujuLogLine, Mount, PeerRelation, Port
+from scenario.state import JujuLogLine, Mount, PeerRelation, Port, Storage
 
 if TYPE_CHECKING:
     from scenario.context import Context
@@ -382,21 +383,46 @@ def action_get(self):
             )
         return action.params
 
-    # TODO:
-    def storage_add(self, *args, **kwargs):  # noqa: U100
-        raise NotImplementedError("storage_add")
+    def storage_add(self, name: str, count: int = 1):
+        if "/" in name:
+            raise ModelError('storage name cannot contain "/"')
 
-    def resource_get(self, *args, **kwargs):  # noqa: U100
-        raise NotImplementedError("resource_get")
+        self._context.requested_storages[name] = count
 
-    def storage_list(self, *args, **kwargs):  # noqa: U100
-        raise NotImplementedError("storage_list")
+    def storage_list(self, name: str) -> List[int]:
+        return [
+            storage.index for storage in self._state.storage if storage.name == name
+        ]
 
-    def storage_get(self, *args, **kwargs):  # noqa: U100
-        raise NotImplementedError("storage_get")
+    def storage_get(self, storage_name_id: str, attribute: str) -> str:
+        if attribute == "location":
+            name, index = storage_name_id.split("/")
+            index = int(index)
+            try:
+                storage: Storage = next(
+                    filter(
+                        lambda s: s.name == name and s.index == index,
+                        self._state.storage,
+                    ),
+                )
+            except StopIteration as e:
+                raise RuntimeError(
+                    f"Storage not found with name={name} and index={index}.",
+                ) from e
+
+            fs_path = storage.get_filesystem(self._context)
+            return str(fs_path)
+
+        raise NotImplementedError(
+            f"storage-get not implemented for attribute={attribute}",
+        )
+
+    def planned_units(self) -> int:
+        return self._state.planned_units
 
-    def planned_units(self, *args, **kwargs):  # noqa: U100
-        raise NotImplementedError("planned_units")
+    # TODO:
+    def resource_get(self, *args, **kwargs):  # noqa: U100
+        raise NotImplementedError("resource_get")
 
 
 class _MockPebbleClient(_TestingPebbleClient):

scenario/runtime.py

@@ -253,6 +253,9 @@ def _get_event_env(self, state: "State", event: "Event", charm_root: Path):
         if container := event.container:
             env.update({"JUJU_WORKLOAD_NAME": container.name})
 
+        if storage := event.storage:
+            env.update({"JUJU_STORAGE_ID": f"{storage.name}/{storage.index}"})
+
         if secret := event.secret:
             env.update(
                 {

scenario/state.py

@@ -646,7 +646,7 @@ def __eq__(self, other):
                 "Comparing Status with Tuples is deprecated and will be removed soon.",
             )
             return (self.name, self.message) == other
-        if isinstance(other, StatusBase):
+        if isinstance(other, (StatusBase, _EntityStatus)):
             return (self.name, self.message) == (other.name, other.message)
         logger.warning(
             f"Comparing Status with {other} is not stable and will be forbidden soon."
@@ -659,12 +659,21 @@ def __iter__(self):
 
     def __repr__(self):
         status_type_name = self.name.title() + "Status"
+        if self.name == "unknown":
+            return f"{status_type_name}()"
         return f"{status_type_name}('{self.message}')"
 
 
 def _status_to_entitystatus(obj: StatusBase) -> _EntityStatus:
     """Convert StatusBase to _EntityStatus."""
-    return _EntityStatus(obj.name, obj.message)
+    statusbase_subclass = type(StatusBase.from_name(obj.name, obj.message))
+
+    class _MyClass(_EntityStatus, statusbase_subclass):
+        # Custom type inheriting from a specific StatusBase subclass to support instance checks:
+        #  isinstance(state.unit_status, ops.ActiveStatus)
+        pass
+
+    return _MyClass(obj.name, obj.message)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -709,6 +718,52 @@ def __post_init__(self):
             )
 
 
+_next_storage_index_counter = 0  # storage indices start at 0
+
+
+def next_storage_index(update=True):
+    """Get the index (used to be called ID) the next Storage to be created will get.
+
+    Pass update=False if you're only inspecting it.
+    Pass update=True if you also want to bump it.
+    """
+    global _next_storage_index_counter
+    cur = _next_storage_index_counter
+    if update:
+        _next_storage_index_counter += 1
+    return cur
+
+
+@dataclasses.dataclass(frozen=True)
+class Storage(_DCBase):
+    """Represents an (attached!) storage made available to the charm container."""
+
+    name: str
+
+    index: int = dataclasses.field(default_factory=next_storage_index)
+    # Every new Storage instance gets a new one, if there's trouble, override.
+
+    def get_filesystem(self, ctx: "Context") -> Path:
+        """Simulated filesystem root in this context."""
+        return ctx._get_storage_root(self.name, self.index)
+
+    @property
+    def attached_event(self) -> "Event":
+        """Sugar to generate a <this storage>-storage-attached event."""
+        return Event(
+            path=normalize_name(self.name + "-storage-attached"),
+            storage=self,
+        )
+
+    @property
+    def detached_event(self) -> "Event":
+        """Sugar to generate a <this storage>-storage-detached event."""
+        return Event(
+            path=normalize_name(self.name + "-storage-detached"),
+            storage=self,
+        )
+
+
 @dataclasses.dataclass(frozen=True)
 class State(_DCBase):
     """Represents the juju-owned portion of a unit's state.
@@ -721,30 +776,48 @@ class State(_DCBase):
     config: Dict[str, Union[str, int, float, bool]] = dataclasses.field(
         default_factory=dict,
     )
+    """The present configuration of this charm."""
     relations: List["AnyRelation"] = dataclasses.field(default_factory=list)
+    """All relations that currently exist for this charm."""
     networks: List[Network] = dataclasses.field(default_factory=list)
+    """All networks currently provisioned for this charm."""
     containers: List[Container] = dataclasses.field(default_factory=list)
+    """All containers (whether they can connect or not) that this charm is aware of."""
+    storage: List[Storage] = dataclasses.field(default_factory=list)
+    """All ATTACHED storage instances for this charm.
+    If a storage is not attached, omit it from this listing."""
 
     # we don't use sets to make json serialization easier
     opened_ports: List[Port] = dataclasses.field(default_factory=list)
+    """Ports opened by juju on this charm."""
     leader: bool = False
+    """Whether this charm has leadership."""
     model: Model = Model()
+    """The model this charm lives in."""
     secrets: List[Secret] = dataclasses.field(default_factory=list)
+    """The secrets this charm has access to (as an owner, or as a grantee)."""
 
+    planned_units: int = 1
+    """Number of non-dying planned units that are expected to be running this application.
+    Use with caution."""
     unit_id: int = 0
+    """ID of the unit hosting this charm."""
     # represents the OF's event queue. These events will be emitted before the event being
     # dispatched, and represent the events that had been deferred during the previous run.
     # If the charm defers any events during "this execution", they will be appended
     # to this list.
     deferred: List["DeferredEvent"] = dataclasses.field(default_factory=list)
+    """Events that have been deferred on this charm by some previous execution."""
     stored_state: List["StoredState"] = dataclasses.field(default_factory=list)
-
-    """Represents the 'juju statuses' of the application/unit being tested."""
+    """Contents of a charm's stored state."""
 
     # the current statuses. Will be cast to _EntitiyStatus in __post_init__
     app_status: Union[StatusBase, _EntityStatus] = _EntityStatus("unknown")
+    """Status of the application."""
     unit_status: Union[StatusBase, _EntityStatus] = _EntityStatus("unknown")
+    """Status of the unit."""
     workload_version: str = ""
+    """Workload version."""
 
     def __post_init__(self):
         for name in ["app_status", "unit_status"]:
@@ -897,6 +970,8 @@ class Event(_DCBase):
     args: Tuple[Any] = ()
     kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
 
+    # if this is a storage event, the storage it refers to
+    storage: Optional["Storage"] = None
     # if this is a relation event, the relation it refers to
     relation: Optional["AnyRelation"] = None
     # and the name of the remote unit this relation event is about