feat: typed testing relation getters

testing/src/scenario/state.py

@@ -69,6 +69,7 @@ class _StateKwargs(TypedDict, total=False):
 UnitID = int
 
 CharmType = TypeVar('CharmType', bound=CharmBase)
+_RelationType = TypeVar('_RelationType', bound='RelationBase')
 
 logger = scenario_logger.getChild('state')
 
@@ -1746,6 +1747,38 @@ def get_relation(self, relation: int, /) -> RelationBase:
                 return state_relation
         raise KeyError(f'relation: id={relation} not found in the State')
 
+    def get_regular_relation(self, relation: int, /) -> Relation:
+        """Get a regular relation from this State by relation id.
+
+        Raises:
+            TypeError: If this relation is not a ``Relation``.
+        """
+        return self._get_typed_relation(relation, kind=Relation)
+
+    def get_peer_relation(self, relation: int, /) -> PeerRelation:
+        """Get a peer relation from this State by relation id.
+
+        Raises:
+            TypeError: If this relation is not a ``PeerRelation``.
+        """
+        return self._get_typed_relation(relation, kind=PeerRelation)
+
+    def get_subordinate_relation(self, relation: int, /) -> SubordinateRelation:
+        """Get a subordinate relation from this State by relation id.
+
+        Raises:
+            TypeError: If this relation is not a ``SubordinateRelation``.
+        """
+        return self._get_typed_relation(relation, kind=SubordinateRelation)
+
+    def _get_typed_relation(self, relation: int, kind: type[_RelationType]) -> _RelationType:
+        rel = self.get_relation(relation)
+        if not isinstance(rel, kind):
+            raise TypeError(
+                f"Relation {relation} is not a {kind.__name__}, it's a {rel.__class__.__name__}"
+            )
+        return rel
+
     def get_relations(self, endpoint: str) -> tuple[RelationBase, ...]:
         """Get all relations on this endpoint from the current state."""
         # we rather normalize the endpoint than worry about cursed metadata situations such as:
@@ -1758,6 +1791,45 @@ def get_relations(self, endpoint: str) -> tuple[RelationBase, ...]:
             r for r in self.relations if _normalise_name(r.endpoint) == normalized_endpoint
         )
 
+    def get_regular_relations(self, endpoint: str) -> tuple[Relation, ...]:
+        """Get regular relations on this endpoint from the current state.
+
+        Raises:
+            TypeError: If any of the relations on this endpoint is not
+                a ``Relation``.
+        """
+        return self._get_typed_relations(endpoint, kind=Relation)
+
+    def get_peer_relations(self, endpoint: str) -> tuple[PeerRelation, ...]:
+        """Get peer relations on this endpoint from the current state.
+
+        Raises:
+            TypeError: If any of the relations on this endpoint is not
+                a ``PeerRelation``.
+        """
+        return self._get_typed_relations(endpoint, kind=PeerRelation)
+
+    def get_subordinate_relations(self, endpoint: str) -> tuple[SubordinateRelation, ...]:
+        """Get subordinate relations on this endpoint from the current state.
+
+        Raises:
+            TypeError: If any of the relations on this endpoint is not
+                a ``SubordinateRelation``.
+        """
+        return self._get_typed_relations(endpoint, kind=SubordinateRelation)
+
+    def _get_typed_relations(
+        self, endpoint: str, kind: type[_RelationType]
+    ) -> tuple[_RelationType, ...]:
+        rels = self.get_relations(endpoint)
+        for rel in rels:
+            if not isinstance(rel, kind):
+                raise TypeError(
+                    f'Relation on endpoint {endpoint} is not a {kind.__name__}'
+                    f", it's a {rel.__class__.__name__}"
+                )
+        return rels
+
     @classmethod
     def from_context(
         cls,

testing/tests/test_e2e/test_relations.py

@@ -672,6 +672,117 @@ def _on_config_changed(self, _):
     assert isinstance(exc.value, KeyError) or isinstance(exc.value.__cause__, KeyError)
 
 
+def test_get_regular_relation():
+    rel = Relation('foo')
+    state = State(relations={rel})
+    rel_out = state.get_regular_relation(rel.id)
+    assert rel_out == rel
+
+
+def test_get_regular_relation_type_error():
+    rel = PeerRelation('foo')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_regular_relation(rel.id)
+
+
+def test_get_peer_relation():
+    rel = PeerRelation('peers')
+    state = State(relations={rel})
+    rel_out = state.get_peer_relation(rel.id)
+    assert rel_out == rel
+
+
+def test_get_peer_relation_type_error():
+    rel = Relation('foo')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_peer_relation(rel.id)
+
+
+def test_get_subordinate_relation():
+    rel = SubordinateRelation('logging')
+    state = State(relations={rel})
+    rel_out = state.get_subordinate_relation(rel.id)
+    assert rel_out == rel
+
+
+def test_get_subordinate_relation_type_error():
+    rel = PeerRelation('peers')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_subordinate_relation(rel.id)
+
+
+def test_get_regular_relations():
+    rel1 = Relation('foo')
+    rel2 = Relation('foo')
+    state = State(relations={rel1, rel2})
+    rels_out = state.get_regular_relations('foo')
+    assert len(rels_out) == 2
+    assert rel1 in rels_out
+    assert rel2 in rels_out
+
+
+def test_get_regular_relations_type_error():
+    rel = PeerRelation('peers')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_regular_relations('peers')
+
+
+def test_get_peer_relations():
+    rel1 = PeerRelation('peers')
+    rel2 = PeerRelation('peers')
+    state = State(relations={rel1, rel2})
+    rels_out = state.get_peer_relations('peers')
+    assert len(rels_out) == 2
+    assert rel1 in rels_out
+    assert rel2 in rels_out
+
+
+def test_get_peer_relations_type_error():
+    rel = SubordinateRelation('logging')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_peer_relations('logging')
+
+
+def test_get_subordinate_relations():
+    rel1 = SubordinateRelation('logging')
+    rel2 = SubordinateRelation('logging')
+    state = State(relations={rel1, rel2})
+    rels_out = state.get_subordinate_relations('logging')
+    assert len(rels_out) == 2
+    assert rel1 in rels_out
+    assert rel2 in rels_out
+
+
+def test_get_subordinate_relations_type_error():
+    rel = Relation('foo')
+    state = State(relations={rel})
+    with pytest.raises(TypeError):
+        state.get_subordinate_relations('foo')
+
+
+def test_get_regular_relations_empty():
+    state = State()
+    rels_out = state.get_regular_relations('foo')
+    assert len(rels_out) == 0
+
+
+def test_get_peer_relations_empty():
+    state = State()
+    rels_out = state.get_peer_relations('peers')
+    assert len(rels_out) == 0
+
+
+def test_get_subordinate_relations_empty():
+    state = State()
+    rels_out = state.get_subordinate_relations('logging')
+    assert len(rels_out) == 0
+
+
 @pytest.mark.parametrize('klass', (Relation, PeerRelation, SubordinateRelation))
 def test_relation_positional_arguments(klass):
     with pytest.raises(TypeError):