diff --git a/docs/reference.md b/docs/reference.md index 83ed603..7b5afdc 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -304,6 +304,27 @@ It is effectively a shortcut for: await ops_test.model.set_config({"update-status-hook-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. diff --git a/pytest_operator/plugin.py b/pytest_operator/plugin.py index 4f5f24c..1946c8a 100644 --- a/pytest_operator/plugin.py +++ b/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": diff --git a/pytest_operator/relation_data.py b/pytest_operator/relation_data.py new file mode 100644 index 0000000..9f5c629 --- /dev/null +++ b/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) diff --git a/tests/data/charms/operator-framework/metadata.yaml b/tests/data/charms/operator-framework/metadata.yaml index c44c8ad..d7eeb07 100644 --- a/tests/data/charms/operator-framework/metadata.yaml +++ b/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 diff --git a/tests/data/charms/operator-framework/src/charm.py b/tests/data/charms/operator-framework/src/charm.py index 1c8f5f9..28ead7c 100755 --- a/tests/data/charms/operator-framework/src/charm.py +++ b/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__": diff --git a/tests/data/charms/reactive-framework/metadata.yaml b/tests/data/charms/reactive-framework/metadata.yaml index 85a6142..7ac6d5d 100644 --- a/tests/data/charms/reactive-framework/metadata.yaml +++ b/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 diff --git a/tests/data/charms/reactive-framework/reactive/__init__.py b/tests/data/charms/reactive-framework/reactive/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/charms/reactive-framework/reactive/reactive_framework.py b/tests/data/charms/reactive-framework/reactive/reactive_framework.py index 62e448b..71f5abe 100644 --- a/tests/data/charms/reactive-framework/reactive/reactive_framework.py +++ b/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 diff --git a/tests/data/charms/reactive-framework/reactive/relations/__init__.py b/tests/data/charms/reactive-framework/reactive/relations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/charms/reactive-framework/reactive/relations/unique-test-relation/__init__.py b/tests/data/charms/reactive-framework/reactive/relations/unique-test-relation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/charms/reactive-framework/reactive/relations/unique-test-relation/requires.py b/tests/data/charms/reactive-framework/reactive/relations/unique-test-relation/requires.py new file mode 100644 index 0000000..602ff6d --- /dev/null +++ b/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] + diff --git a/tests/integration/test_pytest_operator.py b/tests/integration/test_pytest_operator.py index 9650d55..a34b9f1 100644 --- a/tests/integration/test_pytest_operator.py +++ b/tests/integration/test_pytest_operator.py @@ -31,7 +31,7 @@ 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"] }} @@ -39,6 +39,8 @@ async def test_build_and_deploy(self, ops_test): 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 diff --git a/tests/unit/test_pytest_operator.py b/tests/unit/test_pytest_operator.py index 1ec0660..e0dc20c 100644 --- a/tests/unit/test_pytest_operator.py +++ b/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(