Allow Harness.get_relation_data to accept instances (canonical#465)

* Allow Harness.get_relation_data to accept instances

It's tripped me up a few times that you have to pass in the name of the
app or unit rather than the instance itself that you use everywhere
else, so this allows get_relation_data to handle either.

* Update to fix test suite with latest changes to operator framework

Author: johnsca

ops/charm.py

@@ -895,6 +895,7 @@ class ContainerMeta:
         name: Name of container (key in the YAML)
         mounts: :class:`ContainerStorageMeta` mounts available to the container
     """
+
     def __init__(self, name, raw):
         self.name = name
         self._mounts = {}
@@ -956,6 +957,7 @@ class ContainerStorageMeta:
     `location` will not be an accessible attribute, as it would not be possible to determine
     which mount point was desired, and `locations` should be iterated over.
     """
+
     def __init__(self, storage, location):
         self.storage = storage
         self._locations = [location]

ops/model.py

@@ -296,7 +296,7 @@ def __init__(self, name, meta, backend, cache):
         self._is_our_unit = self.name == self._backend.unit_name
         self._status = None
 
-        if self._is_our_unit:
+        if self._is_our_unit and hasattr(meta, "containers"):
             self._containers = ContainerMapping(meta.containers, backend)
 
     def _invalidate(self):
@@ -1052,6 +1052,7 @@ class Container:
     Attributes:
         name: The name of the container from metadata.yaml (eg, 'postgres').
     """
+
     def __init__(self, name, backend, pebble_client=None):
         self.name = name
 

ops/testing.py

@@ -37,6 +37,13 @@
 # pass in a file-like object or the string directly.
 OptionalYAML = typing.Optional[typing.Union[str, typing.TextIO]]
 
+
+# An instance of an Application or Unit, or the name of either.
+# This is done here to avoid a scoping issue with the `model` property
+# of the Harness class below.
+AppUnitOrName = typing.Union[str, model.Application, model.Unit]
+
+
 # CharmType represents user charms that are derived from CharmBase.
 CharmType = typing.TypeVar('CharmType', bound=charm.CharmBase)
 
@@ -72,6 +79,7 @@ class Harness(typing.Generic[CharmType]):
             config.yaml. If not supplied, we will look for a 'config.yaml' file in the
             parent directory of the Charm.
     """
+
     def __init__(
             self,
             charm_cls: typing.Type[CharmType],
@@ -578,7 +586,7 @@ def _emit_relation_departed(self, relation_id, unit_name):
             raise ValueError('Invalid Unit Name')
         self._charm.on[rel_name].relation_departed.emit(relation, app, unit)
 
-    def get_relation_data(self, relation_id: int, app_or_unit: str) -> typing.Mapping:
+    def get_relation_data(self, relation_id: int, app_or_unit: AppUnitOrName) -> typing.Mapping:
         """Get the relation data bucket for a single app or unit in a given relation.
 
         This ignores all of the safety checks of who can and can't see data in relations (eg,
@@ -587,13 +595,15 @@ def get_relation_data(self, relation_id: int, app_or_unit: str) -> typing.Mappin
 
         Args:
             relation_id: The relation whose content we want to look at.
-            app_or_unit: The name of the application or unit whose data we want to read
+            app_or_unit: An Application or Unit instance, or its name, whose data we want to read
         Return:
             A dict containing the relation data for `app_or_unit` or None.
 
         Raises:
             KeyError: if relation_id doesn't exist
         """
+        if hasattr(app_or_unit, 'name'):
+            app_or_unit = app_or_unit.name
         return self._backend._relation_data[relation_id].get(app_or_unit, None)
 
     def get_pod_spec(self) -> (typing.Mapping, typing.Mapping):

test/test_charm.py

@@ -473,8 +473,8 @@ def test_containers_storage_multiple_mounts(self):
         self.assertIsInstance(meta.containers['test1'], ContainerMeta)
         self.assertIsInstance(meta.containers['test1'].mounts["data"], ContainerStorageMeta)
         self.assertEqual(
-                meta.containers['test1'].mounts["data"].locations[0],
-                '/test/storagemount')
+            meta.containers['test1'].mounts["data"].locations[0],
+            '/test/storagemount')
         self.assertEqual(meta.containers['test1'].mounts["data"].locations[1], '/test/otherdata')
 
         with self.assertRaises(RuntimeError):

test/test_pebble.py

@@ -469,13 +469,13 @@ def test_service_equality(self):
                                      raw={
                                           "override": "replace",
                                           "command": "echo foo"
-                                         })
+                                     })
         old_services = {"foo": old_service}
         self.assertEqual(plan.services, old_services)
 
         services_as_dict = {
             "foo": {"override": "replace", "command": "echo foo"}
-            }
+        }
         self.assertEqual(plan.services, services_as_dict)
 
 

test/test_testing.py

@@ -32,17 +32,16 @@
     Object,
 )
 from ops.model import (
+    Application,
     ActiveStatus,
     MaintenanceStatus,
     UnknownStatus,
     ModelError,
     RelationNotFoundError,
     _ModelBackend,
+    Unit,
 )
-from ops.testing import (
-    Harness,
-    _TestingPebbleClient,
-)
+from ops.testing import Harness, _TestingPebbleClient
 
 
 class TestHarness(unittest.TestCase):
@@ -548,12 +547,13 @@ def test_relation_events(self):
               }}])
 
     def test_get_relation_data(self):
-        harness = Harness(CharmBase, meta='''
+        charm_meta = '''
             name: test-app
             requires:
                 db:
                     interface: pgsql
-            ''')
+        '''
+        harness = Harness(CharmBase, meta=charm_meta)
         self.addCleanup(harness.cleanup)
         rel_id = harness.add_relation('db', 'postgresql')
         harness.update_relation_data(rel_id, 'postgresql', {'remote': 'data'})
@@ -565,6 +565,16 @@ def test_get_relation_data(self):
             # unknown relation id
             harness.get_relation_data(99, 'postgresql')
 
+        meta = yaml.safe_load(charm_meta)
+        t_app = Application('test-app', meta, harness._backend, None)
+        t_unit0 = Unit('test-app/0', meta, harness._backend, {Application: t_app})
+        t_unit1 = Unit('test-app/1', meta, harness._backend, {Application: t_app})
+        self.assertEqual(harness.get_relation_data(rel_id, t_app), {})
+        self.assertEqual(harness.get_relation_data(rel_id, t_unit0), {})
+        self.assertEqual(harness.get_relation_data(rel_id, t_unit1), None)
+        pg_app = Application('postgresql', meta, harness._backend, None)
+        self.assertEqual(harness.get_relation_data(rel_id, pg_app), {'remote': 'data'})
+
     def test_create_harness_twice(self):
         metadata = '''
             name: my-charm