Relation data wrapper

docs/reference.md

@@ -304,6 +304,27 @@ It is effectively a shortcut for:
     await ops_test.model.set_config({"update-status-hook-interval": <slow-interval>})
 ```
 
+#### `async def get_relation_data(self, *, provider_endpoint: str, requirer_endpoint: str, include_juju_keys: bool = False, refresh_cache: bool = False) -> RelationData`
+
+Retrieves the relation databags for two given endpoints.
+This can be used to verify that the relation data (unit, app) for either side has
+the expected shape.
+
+Note: Depends on the presense of the juju client snap
+
+`include_juju_keys = True` includes egress-subnets, ingress-address, and private-address
+`refresh_cache = True` will force a read-through cache from the controller
+
+
+```python
+data = await ops_test.get_relation_data(
+    provider_endpoint='my_app/0:ingress',
+    requirer_endpoint='other_app/1:ingress')
+assert data.provider.application_data == {'foo': 'bar', 
+                                          'baz': 'qux'}
+```
+
+
 #### `def abort(self, *args, **kwargs)`
 
 Fail the current test method and mark all remaining test methods as xfail.

pytest_operator/plugin.py

@@ -47,6 +47,8 @@
 from juju.errors import JujuError
 from juju.model import Model, Controller, websockets
 
+from pytest_operator.relation_data import get_relation_data, RelationData
+
 log = logging.getLogger(__name__)
 
 
@@ -1422,6 +1424,39 @@ async def fast_forward(
         yield
         await model.set_config({update_interval_key: interval_after})
 
+    async def get_relation_data(
+        self,
+        *,
+        provider_endpoint: str,
+        requirer_endpoint: str,
+        include_juju_keys: bool = False,
+        refresh_cache: bool = False,
+    ) -> RelationData:
+        """Get relation databag contents for both sides of a juju relation.
+        Note: Depends on the presense of the juju client snap
+
+        include_juju_keys = True includes egress-subnets, ingress-address,
+                                 and private-address
+        refresh_cache = True will force a read-through cache from the controller
+
+        Usage:
+        >>> data = await ops_test.get_relation_data(
+        ...    provider_endpoint='prometheus/0:ingress',
+        ...    requirer_endpoint='traefik/1:ingress-per-unit')
+        >>> assert data.provider.application_data == {'foo': 'bar', 'baz': 'qux'}
+
+        """
+        if not self.model:
+            raise RuntimeError("No model currently set.")
+
+        return await get_relation_data(
+            self.model,
+            provider_endpoint,
+            requirer_endpoint,
+            include_juju_keys,
+            refresh_cache,
+        )
+
     def is_crash_dump_enabled(self) -> bool:
         """Returns whether Juju crash dump is enabled given the current settings."""
         if self.crash_dump == "always":

pytest_operator/relation_data.py

@@ -0,0 +1,183 @@
+import asyncio
+import os
+import shlex
+from dataclasses import dataclass
+from subprocess import PIPE, Popen
+from typing import Dict
+
+from juju.model import Model
+import yaml
+
+_JUJU_DATA_CACHE = {}
+_JUJU_KEYS = ("egress-subnets", "ingress-address", "private-address")
+
+
+def _purge(data: dict):
+    for key in _JUJU_KEYS:
+        if key in data:
+            del data[key]
+
+
+async def _get_unit_info(
+    model: Model, unit_name: str, *, refresh_cache: bool = False
+) -> dict:
+    """Returns unit-info data structure.
+
+     for example:
+
+    traefik-k8s/0:
+      opened-ports: []
+      charm: local:focal/traefik-k8s-1
+      leader: true
+      relation-info:
+      - endpoint: ingress-per-unit
+        related-endpoint: ingress
+        application-data:
+          _supported_versions: '- v1'
+        related-units:
+          prometheus-k8s/0:
+            in-scope: true
+            data:
+              egress-subnets: 10.152.183.150/32
+              ingress-address: 10.152.183.150
+              private-address: 10.152.183.150
+      provider-id: traefik-k8s-0
+      address: 10.1.232.144
+    """
+
+    if cached_data := _JUJU_DATA_CACHE.get(unit_name):
+        if refresh_cache:
+            _JUJU_DATA_CACHE.pop(unit_name)
+        else:
+            return cached_data
+
+    new_env = os.environ.copy()
+    new_env["JUJU_MODEL"] = model.name
+    cmd = shlex.split(f"juju show-unit {unit_name}")
+
+    proc = Popen(cmd, stdout=PIPE, stderr=PIPE, env=new_env)
+    raw_data = proc.stdout.read().decode("utf-8").strip()
+    if not raw_data:
+        raise ValueError(
+            f"no unit info could be grabbed for {unit_name}; "
+            f"are you sure it's a valid unit name?\n"
+            f"{proc.stderr.read().decode()}"
+        )
+
+    data = yaml.safe_load(raw_data)
+    _JUJU_DATA_CACHE[unit_name] = data
+    return data
+
+
+def _get_relation_by_endpoint(relations, endpoint, remote_obj):
+    relations = [
+        r
+        for r in relations
+        if r["endpoint"] == endpoint and remote_obj in r["related-units"]
+    ]
+    if not relations:
+        raise ValueError(f"no relations found with endpoint==" f"{endpoint}")
+    if len(relations) > 1:
+        raise ValueError("multiple relations found with endpoint==" f"{endpoint}")
+    return relations[0]
+
+
+@dataclass
+class UnitRelationData:
+    unit_name: str
+    endpoint: str
+    leader: bool
+    application_data: Dict[str, str]
+    unit_data: Dict[str, str]
+
+
+async def _get_endpoint_content(
+    model: Model,
+    obj: str,
+    other_obj: str,
+    include_default_juju_keys: bool = False,
+    refresh_cache: bool = False,
+) -> UnitRelationData:
+    """Get the content of the databag  `obj` sent to `other_obj`."""
+    endpoint = None
+    other_unit_name = other_obj.split(":")[0] if ":" in other_obj else other_obj
+    if ":" in obj:
+        unit_name, endpoint = obj.split(":")
+    else:
+        unit_name = obj
+    data = (await _get_unit_info(model, unit_name, refresh_cache=refresh_cache))[
+        unit_name
+    ]
+    is_leader = data["leader"]
+
+    relation_infos = data.get("relation-info")
+    if not relation_infos:
+        raise RuntimeError(f"{unit_name} has no relations")
+
+    if not endpoint:
+        relation_data_raw = relation_infos[0]
+        endpoint = relation_data_raw["endpoint"]
+    else:
+        relation_data_raw = _get_relation_by_endpoint(
+            relation_infos, endpoint, other_unit_name
+        )
+
+    related_units_data_raw = relation_data_raw["related-units"]
+
+    if not other_unit_name:
+        other_unit_name = next(iter(related_units_data_raw.keys()))
+    other_unit_info = await _get_unit_info(
+        model, other_unit_name, refresh_cache=refresh_cache
+    )
+    other_unit_relation_infos = other_unit_info[other_unit_name]["relation-info"]
+    remote_data_raw = _get_relation_by_endpoint(
+        other_unit_relation_infos, relation_data_raw["related-endpoint"], unit_name
+    )
+    this_unit_data = remote_data_raw["related-units"][unit_name]["data"]
+    this_app_data = remote_data_raw["application-data"]
+
+    if not include_default_juju_keys:
+        _purge(this_unit_data)
+
+    return UnitRelationData(
+        unit_name, endpoint, is_leader, this_app_data, this_unit_data
+    )
+
+
+@dataclass
+class RelationData:
+    provider: UnitRelationData
+    requirer: UnitRelationData
+
+
+async def get_relation_data(
+    model: Model,
+    provider_endpoint: str,
+    requirer_endpoint: str,
+    include_juju_keys: bool = False,
+    refresh_cache: bool = False,
+) -> RelationData:
+    """Get relation databag contents for both sides of a juju relation.
+
+    Usage:
+    >>> data: RelationData = await ops_test.get_relation_data(
+    ...     'prometheus/0:ingress', 'traefik/1:ingress-per-unit')
+    >>> assert data.provider.application_data['key'] = 'foo'
+    """
+    provider_data, requirer_data = await asyncio.gather(
+        _get_endpoint_content(
+            model,
+            provider_endpoint,
+            requirer_endpoint,
+            include_juju_keys,
+            refresh_cache,
+        ),
+        _get_endpoint_content(
+            model,
+            requirer_endpoint,
+            provider_endpoint,
+            include_juju_keys,
+            refresh_cache,
+        ),
+    )
+    return RelationData(provider_data, requirer_data)

tests/data/charms/operator-framework/metadata.yaml

@@ -5,4 +5,7 @@ description: |
   TODO: fill out the charm's description
 summary: |
   TODO: fill out the charm's summary
-series: [focal]
+series: [focal, jammy]
+provides:
+  operator:
+    interface: unique-test-relation

tests/data/charms/operator-framework/src/charm.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import json
 
 from ops.charm import CharmBase
 from ops.main import main
@@ -9,6 +10,10 @@ class OperatorFrameworkCharm(CharmBase):
     def __init__(self, *args):
         super().__init__(*args)
         self.unit.status = ActiveStatus()
+        self.framework.observe(self.on.operator_relation_joined, self.on_relation_joined)
+
+    def on_relation_joined(self, event):
+        event.relation.data[self.model.unit]["units"] = json.dumps([_.name for _ in event.relation.units])
 
 
 if __name__ == "__main__":

tests/data/charms/reactive-framework/metadata.yaml

@@ -5,5 +5,9 @@ description: |
     a test charm
 series:
   - focal
+  - jammy
 tags:
   - misc
+requires:
+  reactive:
+    interface: unique-test-relation

tests/data/charms/reactive-framework/reactive/reactive_framework.py

@@ -1,4 +1,4 @@
-from charms.reactive import when_not
+from charms.reactive import  when_not
 from charms import layer
 
 

tests/data/charms/reactive-framework/reactive/relations/unique-test-relation/requires.py

@@ -0,0 +1,9 @@
+from charms.reactive import when, Endpoint
+
+
+class UniqueTestRelationRequires(Endpoint):
+    @when('endpoint.{endpoint_name}.joined')
+    def joined(self):
+        for relation in self.relations:
+            relation.to_publish["units"] = [_.unit_name for _ in relation.joined_units]
+

tests/integration/test_pytest_operator.py

@@ -31,14 +31,16 @@ async def test_build_and_deploy(self, ops_test):
             # test as an example, I included it directly here, since it's small. E.g.:
             # "tests/data/bundle.yaml",
             """
-                series: focal
+                series: jammy
                 applications:
                   reactive-framework:
                     charm: {{ charms["reactive-framework"] }}
                     num_units: 1
                   operator-framework:
                     charm: {{ charms["operator-framework"] }}
                     num_units: 1
+                relations:
+                    - ["reactive-framework:reactive", "operator-framework:operator"]
             """,
             charms=await ops_test.build_charms(*charms),
         )
@@ -105,6 +107,25 @@ async def test_3_context_failure_reverts_model(self, ops_test):
                 raise ZeroDivisionError()
         assert ops_test.current_alias == prior_alias
 
+    async def test_4_get_relation_data(self, ops_test):
+        related = await ops_test.get_relation_data(
+            provider_endpoint="operator-framework/0:operator",
+            requirer_endpoint="reactive-framework/0:reactive",
+        )
+        assert related.provider.unit_data["units"] == '["reactive-framework/0"]'
+        assert related.requirer.unit_data["units"] == '["operator-framework/0"]'
+        assert related.provider.leader, "Should be leader"
+        assert related.requirer.leader, "Should be leader"
+
+        related = await ops_test.get_relation_data(
+            provider_endpoint="operator-framework/0:operator",
+            requirer_endpoint="reactive-framework/0",
+        )
+        assert related.provider.unit_data["units"] == '["reactive-framework/0"]'
+        assert related.requirer.unit_data["units"] == '["operator-framework/0"]'
+        assert related.provider.leader, "Should be leader"
+        assert related.requirer.leader, "Should be leader"
+
 
 async def test_func(ops_test):
     assert ops_test.model

tests/unit/test_pytest_operator.py

@@ -276,7 +276,7 @@ async def test_plugin_fetch_resources(tmp_path_factory, resource_charm):
     arch_resources = ops_test.arch_specific_resources(resource_charm)
 
     def dl_rsc(resource, dest_path):
-        assert type(resource) == str
+        assert isinstance(resource, str)
         return dest_path
 
     with patch(