Fix setting value of relationship of type pool (#118)

Fixes #27

Author: gmazoyer

changelog/27.fixed.md

@@ -0,0 +1 @@
+Fix generated GraphQL query when having a relationship to a pool node

infrahub_sdk/node.py

@@ -241,10 +241,10 @@ def typename(self) -> Optional[str]:
             return self._peer.typename
         return self._typename
 
-    def _generate_input_data(self) -> dict[str, Any]:
+    def _generate_input_data(self, allocate_from_pool: bool = False) -> dict[str, Any]:
         data: dict[str, Any] = {}
 
-        if self.is_resource_pool:
+        if self.is_resource_pool and allocate_from_pool:
             return {"from_pool": {"id": self.id}}
 
         if self.id is not None:
@@ -424,8 +424,8 @@ def peer_hfids_str(self) -> list[str]:
     def has_update(self) -> bool:
         return self._has_update
 
-    def _generate_input_data(self) -> list[dict]:
-        return [peer._generate_input_data() for peer in self.peers]
+    def _generate_input_data(self, allocate_from_pool: bool = False) -> list[dict]:
+        return [peer._generate_input_data(allocate_from_pool=allocate_from_pool) for peer in self.peers]
 
     def _generate_mutation_query(self) -> dict[str, Any]:
         # Does nothing for now
@@ -818,6 +818,7 @@ def _generate_input_data(self, exclude_unmodified: bool = False, exclude_hfid: b
                 data[item_name] = attr_data
 
         for item_name in self._relationships:
+            allocate_from_pool = False
             rel_schema = self._schema.get_relationship(name=item_name)
             if not rel_schema or rel_schema.read_only:
                 continue
@@ -836,7 +837,12 @@ def _generate_input_data(self, exclude_unmodified: bool = False, exclude_hfid: b
             if rel is None or not rel.initialized:
                 continue
 
-            rel_data = rel._generate_input_data()
+            if isinstance(rel, (RelatedNode, RelatedNodeSync)) and rel.is_resource_pool:
+                # If the relatiionship is a resource pool and the expected schema is different from the one of the pool, this means we expect to get
+                # a resource from the pool itself
+                allocate_from_pool = rel_schema.peer != rel.peer._schema.kind
+
+            rel_data = rel._generate_input_data(allocate_from_pool=allocate_from_pool)
 
             if rel_data and isinstance(rel_data, dict):
                 if variable_values := rel_data.get("data"):
@@ -1426,7 +1432,7 @@ async def get_pool_allocated_resources(self, resource: InfrahubNode) -> list[Inf
             list[InfrahubNode]: The allocated nodes.
         """
         if not self.is_resource_pool():
-            raise ValueError("Allocate resources can only be fetched from resource pool nodes.")
+            raise ValueError("Allocated resources can only be fetched from resource pool nodes.")
 
         graphql_query_name = "InfrahubResourcePoolAllocated"
         node_ids_per_kind: dict[str, list[str]] = {}

tests/unit/sdk/conftest.py

@@ -821,7 +821,15 @@ async def simple_device_schema() -> NodeSchema:
                 "optional": True,
                 "cardinality": "one",
                 "kind": "Attribute",
-            }
+            },
+            {
+                "name": "ip_address_pool",
+                "peer": "CoreIPAddressPool",
+                "label": "Address allocator",
+                "optional": True,
+                "cardinality": "one",
+                "kind": "Attribute",
+            },
         ],
     }
     return NodeSchema(**data)  # type: ignore

tests/unit/sdk/test_node.py

@@ -1487,7 +1487,9 @@ async def test_create_input_data_with_resource_pool_relationship(
             },
         )
         device = InfrahubNode(
-            client=client, schema=simple_device_schema, data={"name": "device-01", "primary_address": ip_pool}
+            client=client,
+            schema=simple_device_schema,
+            data={"name": "device-01", "primary_address": ip_pool, "ip_address_pool": ip_pool},
         )
     else:
         ip_prefix = InfrahubNodeSync(client=client, schema=ipam_ipprefix_schema, data=ipam_ipprefix_data)
@@ -1504,13 +1506,16 @@ async def test_create_input_data_with_resource_pool_relationship(
             },
         )
         device = InfrahubNode(
-            client=client, schema=simple_device_schema, data={"name": "device-01", "primary_address": ip_pool}
+            client=client,
+            schema=simple_device_schema,
+            data={"name": "device-01", "primary_address": ip_pool, "ip_address_pool": ip_pool},
         )
 
     assert device._generate_input_data()["data"] == {
         "data": {
             "name": {"value": "device-01"},
             "primary_address": {"from_pool": {"id": "pppppppp-pppp-pppp-pppp-pppppppppppp"}},
+            "ip_address_pool": {"id": "pppppppp-pppp-pppp-pppp-pppppppppppp"},
         },
     }
 
@@ -1534,7 +1539,9 @@ async def test_create_mutation_query_with_resource_pool_relationship(
             },
         )
         device = InfrahubNode(
-            client=client, schema=simple_device_schema, data={"name": "device-01", "primary_address": ip_pool}
+            client=client,
+            schema=simple_device_schema,
+            data={"name": "device-01", "primary_address": ip_pool, "ip_address_pool": ip_pool},
         )
     else:
         ip_prefix = InfrahubNodeSync(client=client, schema=ipam_ipprefix_schema, data=ipam_ipprefix_data)
@@ -1551,11 +1558,17 @@ async def test_create_mutation_query_with_resource_pool_relationship(
             },
         )
         device = InfrahubNode(
-            client=client, schema=simple_device_schema, data={"name": "device-01", "primary_address": ip_pool}
+            client=client,
+            schema=simple_device_schema,
+            data={"name": "device-01", "primary_address": ip_pool, "ip_address_pool": ip_pool},
         )
 
     assert device._generate_mutation_query() == {
-        "object": {"id": None, "primary_address": {"node": {"__typename": None, "display_label": None, "id": None}}},
+        "object": {
+            "id": None,
+            "primary_address": {"node": {"__typename": None, "display_label": None, "id": None}},
+            "ip_address_pool": {"node": {"__typename": None, "display_label": None, "id": None}},
+        },
         "ok": None,
     }