Repository: canonical/operator

Issue:
Testing harness: cannot find StoredVariable attribute in type <CharmClass>
When the testing harness is used, having a state attribute on a charm [currently](https://github.com/canonical/operator/commit/44dff930667aa8e9b179c11fa87ceb8c9b85ec5a) raises an exception.

This happens because `Harness.begin` creates a new type by inheriting from the real charm class provided by a test developer so when the `__dict__` is used on it in the `StoredState.__get__`, the attribute lookup fails:
https://github.com/canonical/operator/blob/44dff930667aa8e9b179c11fa87ceb8c9b85ec5a/ops/testing.py#L95-L96

The lookup in `parent_type.__bases__[0].__dict__.items()` would succeed.

Example:
https://github.com/dshcherb/charm-haproxy/blob/ba9bae5aae3b89772361bfd4a973bb61d5abca83/test/test_charm.py#L21

```python3 -m unittest
[6] > /home/ubuntu/src/canonical/juju/charm-haproxy/lib/ops/framework.py(834)__get__()
-> if bound is None:
   13 frames hidden (try 'help hidden_frames')
(Pdb++) l
829  	            return self
830  	
831  	        bound = parent.__dict__.get(self.attr_name)
832  	        import pdb
833  	        pdb.set_trace()
834  ->	        if bound is None:
835  	            for attr_name, attr_value in parent_type.__dict__.items():
836  	                if attr_value is self:
837  	                    if self.attr_name and attr_name != self.attr_name:
838  	                        parent_tname = parent_type.__name__
839  	                        raise RuntimeError("StoredState shared by {}.{} and {}.{}".format(
(Pdb++) parent_type.__dict__.items()
dict_items([('__module__', 'ops.testing'), ('on', <ops.testing.Harness.begin.<locals>.TestEvents object at 0x7f935be611f0>), ('__doc__', None)])

(Pdb++) parent_type.__bases__[0].__dict__.items()
dict_items([('__module__', 'src.charm'), ('state', <ops.framework.StoredState object at 0x7f935c486400>), ('HAPROXY_ENV_FILE', PosixPath('/etc/default/haproxy')), ('__init__', <function HaproxyCharm.__init__ at 0x7f935c06faf0>), ('on_install', <function HaproxyCharm.on_install at 0x7f935befc160>), ('on_start', <function HaproxyCharm.on_start at 0x7f935befc1f0>), ('on_stop', <function HaproxyCharm.on_stop at 0x7f935befc280>), ('on_config_changed', <function HaproxyCharm.on_config_changed at 0x7f935befc310>), ('on_backends_changed', <function HaproxyCharm.on_backends_changed at 0x7f935befc3a0>), ('__doc__', None)])

(Pdb++) c
E
======================================================================
ERROR: test_install (test.test_charm.TestTestingModelBackend)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/lib/python3.8/unittest/mock.py", line 1348, in patched
    return func(*newargs, **newkeywargs)
  File "/home/ubuntu/src/canonical/juju/charm-haproxy/test/test_charm.py", line 21, in test_install
    harness.begin()
  File "lib/ops/testing.py", line 101, in begin
    self._charm = TestCharm(self._framework, self._framework.meta.name)
  File "/home/ubuntu/src/canonical/juju/charm-haproxy/src/charm.py", line 41, in __init__
    self.state.set_default(started=False)
  File "lib/ops/framework.py", line 834, in __get__
    if bound is None:
RuntimeError: cannot find StoredVariable attribute in type HaproxyCharm

----------------------------------------------------------------------
Ran 1 test in 30.287s

FAILED (errors=1)
```

Comments:

Patch:
diff --git a/ops/framework.py b/ops/framework.py
index abdd1fe..c8df98b 100755
--- a/ops/framework.py
+++ b/ops/framework.py
@@ -812,39 +812,78 @@ class BoundStoredState:
 
 
 class StoredState:
+    """A class used to store data the charm needs persisted across invocations.
+
+    Example::
+
+        class MyClass(Object):
+            _stored = StoredState()
+
+    Instances of `MyClass` can transparently save state between invocations by
+    setting attributes on `stored`. Initial state should be set with
+    `set_default` on the bound object, that is::
+
+        class MyClass(Object):
+            _stored = StoredState()
+
+        def __init__(self, parent, key):
+            super().__init__(parent, key)
+            self._stored.set_default(seen=set())
+            self.framework.observe(self.on.seen, self._on_seen)
+
+        def _on_seen(self, event):
+            self._stored.seen.add(event.uuid)
+
+    """
 
     def __init__(self):
         self.parent_type = None
         self.attr_name = None
 
     def __get__(self, parent, parent_type=None):
-        if self.parent_type is None:
-            self.parent_type = parent_type
-        elif self.parent_type is not parent_type:
+        if self.parent_type is not None and self.parent_type not in parent_type.mro():
+            # the StoredState instance is being shared between two unrelated classes
+            # -> unclear what is exepcted of us -> bail out
             raise RuntimeError(
                 'StoredState shared by {} and {}'.format(
                     self.parent_type.__name__, parent_type.__name__))
 
         if parent is None:
+            # accessing via the class directly (e.g. MyClass.stored)
             return self
 
-        bound = parent.__dict__.get(self.attr_name)
-        if bound is None:
-            for attr_name, attr_value in parent_type.__dict__.items():
-                if attr_value is self:
-                    if self.attr_name and attr_name != self.attr_name:
-                        parent_tname = parent_type.__name__
-                        raise RuntimeError("StoredState shared by {}.{} and {}.{}".format(
-                            parent_tname, self.attr_name, parent_tname, attr_name))
-                    self.attr_name = attr_name
-                    bound = BoundStoredState(parent, attr_name)
-                    parent.__dict__[attr_name] = bound
-                    break
-            else:
-                raise RuntimeError(
-                    'cannot find StoredVariable attribute in type {}'.format(parent_type.__name__))
-
-        return bound
+        bound = None
+        if self.attr_name is not None:
+            bound = parent.__dict__.get(self.attr_name)
+            if bound is not None:
+                # we already have the thing from a previous pass, huzzah
+                return bound
+
+        # need to find ourselves amongst the parent's bases
+        for cls in parent_type.mro():
+            for attr_name, attr_value in cls.__dict__.items():
+                if attr_value is not self:
+                    continue
+                # we've found ourselves! is it the first time?
+                if bound is not None:
+                    # the StoredState instance is being stored in two different
+                    # attributes -> unclear what is expected of us -> bail out
+                    raise RuntimeError("StoredState shared by {0}.{1} and {0}.{2}".format(
+                        cls.__name__, self.attr_name, attr_name))
+                # we've found ourselves for the first time; save where, and bind the object
+                self.attr_name = attr_name
+                self.parent_type = cls
+                bound = BoundStoredState(parent, attr_name)
+
+        if bound is not None:
+            # cache the bound object to avoid the expensive lookup the next time
+            # (don't use setattr, to keep things symmetric with the fast-path lookup above)
+            parent.__dict__[self.attr_name] = bound
+            return bound
+
+        raise AttributeError(
+            'cannot find {} attribute in type {}'.format(
+                self.__class__.__name__, parent_type.__name__))
 
 
 def _wrap_stored(parent_data, value):


Test patch:
diff --git a/test/test_framework.py b/test/test_framework.py
index 7671075..1977c1e 100644
--- a/test/test_framework.py
+++ b/test/test_framework.py
@@ -774,12 +774,63 @@ class TestStoredState(unittest.TestCase):
         return framework
 
     def test_basic_state_storage(self):
-        framework = self.create_framework()
+        class SomeObject(Object):
+            state = StoredState()
+
+        self._stored_state_tests(SomeObject)
+
+    def test_straight_subclass(self):
+        class SomeObject(Object):
+            state = StoredState()
+
+        class Sub(SomeObject):
+            pass
 
+        self._stored_state_tests(Sub)
+
+    def test_straight_sub_subclass(self):
         class SomeObject(Object):
             state = StoredState()
 
-        obj = SomeObject(framework, "1")
+        class Sub(SomeObject):
+            pass
+
+        class SubSub(SomeObject):
+            pass
+
+        self._stored_state_tests(SubSub)
+
+    def test_two_subclasses(self):
+        class SomeObject(Object):
+            state = StoredState()
+
+        class SubA(SomeObject):
+            pass
+
+        class SubB(SomeObject):
+            pass
+
+        self._stored_state_tests(SubA)
+        self._stored_state_tests(SubB)
+
+    def test_the_crazy_thing(self):
+        class NoState(Object):
+            pass
+
+        class StatedObject(NoState):
+            state = StoredState()
+
+        class Sibling(NoState):
+            pass
+
+        class FinalChild(StatedObject, Sibling):
+            pass
+
+        self._stored_state_tests(FinalChild)
+
+    def _stored_state_tests(self, cls):
+        framework = self.create_framework()
+        obj = cls(framework, "1")
 
         try:
             obj.state.foo
@@ -812,12 +863,88 @@ class TestStoredState(unittest.TestCase):
 
         # Since this has the same absolute object handle, it will get its state back.
         framework_copy = self.create_framework()
-        obj_copy = SomeObject(framework_copy, "1")
+        obj_copy = cls(framework_copy, "1")
         self.assertEqual(obj_copy.state.foo, 42)
         self.assertEqual(obj_copy.state.bar, "s")
         self.assertEqual(obj_copy.state.baz, 4.2)
         self.assertEqual(obj_copy.state.bing, True)
 
+        framework_copy.close()
+
+    def test_two_subclasses_no_conflicts(self):
+        class Base(Object):
+            state = StoredState()
+
+        class SubA(Base):
+            pass
+
+        class SubB(Base):
+            pass
+
+        framework = self.create_framework()
+        a = SubA(framework, None)
+        b = SubB(framework, None)
+        z = Base(framework, None)
+
+        a.state.foo = 42
+        b.state.foo = "hello"
+        z.state.foo = {1}
+
+        framework.commit()
+        framework.close()
+
+        framework2 = self.create_framework()
+        a2 = SubA(framework2, None)
+        b2 = SubB(framework2, None)
+        z2 = Base(framework2, None)
+
+        self.assertEqual(a2.state.foo, 42)
+        self.assertEqual(b2.state.foo, "hello")
+        self.assertEqual(z2.state.foo, {1})
+
+    def test_two_names_one_state(self):
+        class Mine(Object):
+            state = StoredState()
+            stored = state
+
+        framework = self.create_framework()
+        obj = Mine(framework, None)
+
+        with self.assertRaises(RuntimeError):
+            obj.state.foo = 42
+
+        framework.close()
+
+        # make sure we're not changing the object on failure
+        self.assertNotIn("stored", obj.__dict__)
+        self.assertNotIn("state", obj.__dict__)
+
+    def test_same_name_two_classes(self):
+        class Base(Object):
+            pass
+
+        class A(Base):
+            stored = StoredState()
+
+        class B(Base):
+            stored = A.stored
+
+        framework = self.create_framework()
+        a = A(framework, None)
+        b = B(framework, None)
+
+        # NOTE it's the second one that actually triggers the
+        # exception, but that's an implementation detail
+        a.stored.foo = 42
+
+        with self.assertRaises(RuntimeError):
+            b.stored.foo = "xyzzy"
+
+        framework.close()
+
+        # make sure we're not changing the object on failure
+        self.assertNotIn("stored", b.__dict__)
+
     def test_mutable_types_invalid(self):
         framework = self.create_framework()