Merge pull request #6465 from opsmill/pog-node-trigger-tests

ogenstad · web-flow · commit c306518f057a · 2025-05-14T16:08:57.000+02:00
Additional tests and guard rails for triggers and actions
diff --git a/backend/infrahub/actions/models.py b/backend/infrahub/actions/models.py
@@ -71,7 +71,7 @@ class CoreNodeTriggerAttributeMatch(CoreNodeTriggerMatch):
 class CoreNodeTriggerRelationshipMatch(CoreNodeTriggerMatch):
     relationship_name: str
     added: bool
-    peer: str
+    peer: str | None
 
 
 class CoreNodeTriggerRule(CoreTriggerRule):
@@ -138,9 +138,10 @@ def _from_node_trigger(
                 event_trigger.match_related = {
                     "prefect.resource.role": "infrahub.node.relationship_update",
                     "infrahub.field.name": match.relationship_name,
-                    "infrahub.relationship.peer_id": match.peer,
                     "infrahub.relationship.peer_status": peer_status,
                 }
+                if isinstance(match.peer, str):
+                    event_trigger.match_related["infrahub.relationship.peer_id"] = match.peer
 
         if isinstance(trigger_rule.action, CoreGeneratorAction):
             workflow = ExecuteWorkflow(
diff --git a/backend/infrahub/graphql/manager.py b/backend/infrahub/graphql/manager.py
@@ -25,7 +25,7 @@
 from .directives import DIRECTIVES
 from .enums import generate_graphql_enum, get_enum_attribute_type_name
 from .metrics import SCHEMA_GENERATE_GRAPHQL_METRICS
-from .mutations.action import InfrahubTriggerRuleMutation
+from .mutations.action import InfrahubTriggerRuleMatchMutation, InfrahubTriggerRuleMutation
 from .mutations.artifact_definition import InfrahubArtifactDefinitionMutation
 from .mutations.ipam import (
     InfrahubIPAddressMutation,
@@ -526,6 +526,8 @@ def generate_mutation_mixin(self) -> type[object]:
                 InfrahubKind.STANDARDWEBHOOK: InfrahubWebhookMutation,
                 InfrahubKind.CUSTOMWEBHOOK: InfrahubWebhookMutation,
                 InfrahubKind.NODETRIGGERRULE: InfrahubTriggerRuleMutation,
+                InfrahubKind.NODETRIGGERATTRIBUTEMATCH: InfrahubTriggerRuleMatchMutation,
+                InfrahubKind.NODETRIGGERRELATIONSHIPMATCH: InfrahubTriggerRuleMatchMutation,
             }
 
             if isinstance(node_schema, NodeSchema) and node_schema.is_ip_prefix():
diff --git a/backend/infrahub/graphql/mutations/action.py b/backend/infrahub/graphql/mutations/action.py
@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from graphene import InputObjectType, Mutation
 from typing_extensions import Self
 
-from infrahub.core.schema import NodeSchema
+from infrahub.core.protocols import CoreNodeTriggerAttributeMatch, CoreNodeTriggerRelationshipMatch, CoreNodeTriggerRule
+from infrahub.exceptions import SchemaNotFoundError, ValidationError
 from infrahub.log import get_logger
 
 from .main import InfrahubMutationMixin, InfrahubMutationOptions
@@ -15,8 +16,11 @@
 
     from infrahub.core.branch import Branch
     from infrahub.core.node import Node
+    from infrahub.core.schema import NodeSchema
     from infrahub.database import InfrahubDatabase
 
+    from ..initialization import GraphqlContext
+
 log = get_logger()
 
 
@@ -28,10 +32,6 @@ def __init_subclass_with_meta__(
         _meta: Any | None = None,
         **options: dict[str, Any],
     ) -> None:
-        # Make sure schema is a valid NodeSchema Node Class
-        if not isinstance(schema, NodeSchema):
-            raise ValueError(f"You need to pass a valid NodeSchema in '{cls.__name__}.Meta', received '{schema}'")
-
         if not _meta:
             _meta = InfrahubMutationOptions(cls)
 
@@ -45,9 +45,12 @@ async def mutate_create(
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: InfrahubDatabase | None = None,  # noqa: ARG003
+        database: InfrahubDatabase | None = None,
     ) -> tuple[Node, Self]:
-        trigger_rule_definition, result = await super().mutate_create(info=info, data=data, branch=branch)
+        graphql_context: GraphqlContext = info.context
+        db = database or graphql_context.db
+        _validate_node_kind(data=data, db=db)
+        trigger_rule_definition, result = await super().mutate_create(info=info, data=data, branch=branch, database=db)
 
         return trigger_rule_definition, result
 
@@ -57,9 +60,105 @@ async def mutate_update(
         info: GraphQLResolveInfo,
         data: InputObjectType,
         branch: Branch,
-        database: InfrahubDatabase | None = None,  # noqa: ARG003
+        database: InfrahubDatabase | None = None,
         node: Node | None = None,  # noqa: ARG003
     ) -> tuple[Node, Self]:
-        trigger_rule_definition, result = await super().mutate_update(info=info, data=data, branch=branch)
+        graphql_context: GraphqlContext = info.context
+        db = database or graphql_context.db
+        _validate_node_kind(data=data, db=db)
+        trigger_rule_definition, result = await super().mutate_update(info=info, data=data, branch=branch, database=db)
 
         return trigger_rule_definition, result
+
+
+class InfrahubTriggerRuleMatchMutation(InfrahubMutationMixin, Mutation):
+    @classmethod
+    def __init_subclass_with_meta__(
+        cls,
+        schema: NodeSchema,
+        _meta: Any | None = None,
+        **options: dict[str, Any],
+    ) -> None:
+        if not _meta:
+            _meta = InfrahubMutationOptions(cls)
+
+        _meta.schema = schema
+
+        super().__init_subclass_with_meta__(_meta=_meta, **options)
+
+    @classmethod
+    async def mutate_create(
+        cls,
+        info: GraphQLResolveInfo,
+        data: InputObjectType,
+        branch: Branch,
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
+    ) -> tuple[Node, Self]:
+        graphql_context: GraphqlContext = info.context
+
+        async with graphql_context.db.start_transaction() as dbt:
+            trigger_match, result = await super().mutate_create(info=info, data=data, branch=branch, database=dbt)
+            trigger_match_model = cast(CoreNodeTriggerAttributeMatch | CoreNodeTriggerRelationshipMatch, trigger_match)
+            node_trigger_rule = await trigger_match_model.trigger.get_peer(db=dbt, raise_on_error=True)
+            node_trigger_rule_model = cast(CoreNodeTriggerRule, node_trigger_rule)
+            node_schema = dbt.schema.get_node_schema(name=node_trigger_rule_model.node_kind.value, duplicate=False)
+            _validate_node_kind_field(data=data, node_schema=node_schema)
+
+        return trigger_match, result
+
+    @classmethod
+    async def mutate_update(
+        cls,
+        info: GraphQLResolveInfo,
+        data: InputObjectType,
+        branch: Branch,
+        database: InfrahubDatabase | None = None,  # noqa: ARG003
+        node: Node | None = None,  # noqa: ARG003
+    ) -> tuple[Node, Self]:
+        graphql_context: GraphqlContext = info.context
+        async with graphql_context.db.start_transaction() as dbt:
+            trigger_match, result = await super().mutate_update(info=info, data=data, branch=branch, database=dbt)
+            trigger_match_model = cast(CoreNodeTriggerAttributeMatch | CoreNodeTriggerRelationshipMatch, trigger_match)
+            node_trigger_rule = await trigger_match_model.trigger.get_peer(db=dbt, raise_on_error=True)
+            node_trigger_rule_model = cast(CoreNodeTriggerRule, node_trigger_rule)
+            node_schema = dbt.schema.get_node_schema(name=node_trigger_rule_model.node_kind.value, duplicate=False)
+            _validate_node_kind_field(data=data, node_schema=node_schema)
+
+        return trigger_match, result
+
+
+def _validate_node_kind(data: InputObjectType, db: InfrahubDatabase) -> None:
+    input_data = cast(dict[str, dict[str, Any]], data)
+    if node_kind := input_data.get("node_kind"):
+        value = node_kind.get("value")
+        if isinstance(value, str):
+            try:
+                db.schema.get_node_schema(name=value, duplicate=False)
+            except SchemaNotFoundError as exc:
+                raise ValidationError(
+                    input_value={"node_kind": "The requested node_kind schema was not found"}
+                ) from exc
+            except ValueError as exc:
+                raise ValidationError(input_value={"node_kind": "The requested node_kind is not a valid node"}) from exc
+
+
+def _validate_node_kind_field(data: InputObjectType, node_schema: NodeSchema) -> None:
+    input_data = cast(dict[str, dict[str, Any]], data)
+    if attribute_name := input_data.get("attribute_name"):
+        value = attribute_name.get("value")
+        if isinstance(value, str):
+            if value not in node_schema.attribute_names:
+                raise ValidationError(
+                    input_value={
+                        "attribute_name": f"The attribute {value} doesn't exist on related node trigger using {node_schema.kind}"
+                    }
+                )
+    if relationship_name := input_data.get("relationship_name"):
+        value = relationship_name.get("value")
+        if isinstance(value, str):
+            if value not in node_schema.relationship_names:
+                raise ValidationError(
+                    input_value={
+                        "relationship_name": f"The relationship {value} doesn't exist on related node trigger using {node_schema.kind}"
+                    }
+                )
diff --git a/backend/tests/unit/actions/test_gather.py b/backend/tests/unit/actions/test_gather.py
@@ -9,6 +9,7 @@
     CoreGroupAction,
     CoreGroupTriggerRule,
     CoreNodeTriggerAttributeMatch,
+    CoreNodeTriggerRelationshipMatch,
     CoreNodeTriggerRule,
     CoreRepository,
     CoreStandardGroup,
@@ -144,6 +145,108 @@ async def test_gather_trigger_gather_trigger_action_rules_node_attribute(
     assert len(triggers) == 0
 
 
+async def test_gather_trigger_gather_trigger_action_rules_node_relationship(
+    register_core_models_schema: SchemaBranch, car_person_schema: SchemaBranch, db: InfrahubDatabase
+) -> None:
+    group_action_target = await Node.init(db=db, schema=CoreStandardGroup)
+    await group_action_target.new(db=db, name="GroupActionTarget")
+    await group_action_target.save(db=db)
+
+    group_action = await Node.init(db=db, schema=CoreGroupAction)
+    await group_action.new(db=db, name="MainGroupAction", group=group_action_target)
+    await group_action.save(db=db)
+
+    car_owner = await Node.init(db=db, schema="TestPerson")
+    await car_owner.new(db=db, name="Bobby")
+    await car_owner.save(db=db)
+
+    main_node_trigger_rule = await Node.init(db=db, schema=CoreNodeTriggerRule)
+    await main_node_trigger_rule.new(
+        db=db,
+        name="main_node_trigger",
+        node_kind="TestCar",
+        mutation_action="created",
+        action=group_action,
+        branch_scope="all_branches",
+    )
+    await main_node_trigger_rule.save(db=db)
+
+    triggers = await gather_trigger_action_rules(db=db)
+    assert len(triggers) == 1
+    automation = triggers[0]
+
+    assert automation.trigger == EventTrigger(
+        events={"infrahub.node.created"},
+        match={"infrahub.node.kind": "TestCar"},
+        match_related={},
+    )
+
+    relationship_match = await Node.init(db=db, schema=CoreNodeTriggerRelationshipMatch)
+    await relationship_match.new(
+        db=db,
+        relationship_name="owner",
+        peer=car_owner.id,
+        trigger=main_node_trigger_rule,
+    )
+    await relationship_match.save(db=db)
+
+    triggers = await gather_trigger_action_rules(db=db)
+    assert len(triggers) == 1
+    automation = triggers[0]
+    assert automation.trigger == EventTrigger(
+        events={"infrahub.node.created"},
+        match={"infrahub.node.kind": "TestCar"},
+        match_related={
+            "prefect.resource.role": "infrahub.node.relationship_update",
+            "infrahub.field.name": "owner",
+            "infrahub.relationship.peer_id": car_owner.id,
+            "infrahub.relationship.peer_status": "added",
+        },
+    )
+
+    relationship_match.added.value = False
+    await relationship_match.save(db=db)
+
+    triggers = await gather_trigger_action_rules(db=db)
+    assert len(triggers) == 1
+    automation = triggers[0]
+    assert automation.trigger == EventTrigger(
+        events={"infrahub.node.created"},
+        match={"infrahub.node.kind": "TestCar"},
+        match_related={
+            "prefect.resource.role": "infrahub.node.relationship_update",
+            "infrahub.field.name": "owner",
+            "infrahub.relationship.peer_id": car_owner.id,
+            "infrahub.relationship.peer_status": "removed",
+        },
+    )
+
+    main_node_trigger_rule.branch_scope.value = "other_branches"
+    main_node_trigger_rule.mutation_action.value = "deleted"
+    await main_node_trigger_rule.save(db=db)
+
+    triggers = await gather_trigger_action_rules(db=db)
+    assert len(triggers) == 1
+    automation = triggers[0]
+
+    assert automation.trigger == EventTrigger(
+        events={"infrahub.node.deleted"},
+        match={"infrahub.node.kind": "TestCar", "infrahub.branch.name": "!main"},
+        match_related={
+            "prefect.resource.role": "infrahub.node.relationship_update",
+            "infrahub.field.name": "owner",
+            "infrahub.relationship.peer_id": car_owner.id,
+            "infrahub.relationship.peer_status": "removed",
+        },
+    )
+
+    main_node_trigger_rule.active.value = False
+    await main_node_trigger_rule.save(db=db)
+
+    triggers = await gather_trigger_action_rules(db=db)
+    assert len(triggers) == 0
+
+
 async def test_gather_trigger_gather_trigger_action_rules_group_generators(
     register_core_models_schema: SchemaBranch, db: InfrahubDatabase
 ) -> None:
diff --git a/backend/tests/unit/graphql/mutations/test_action.py b/backend/tests/unit/graphql/mutations/test_action.py
