fix test include

Author: LucasG0

infrahub_sdk/client.py

@@ -558,6 +558,9 @@ async def _process_nodes_and_relationships(
                 - 'related_nodes': A list of InfrahubNode objects representing the related nodes
         """
 
+        # Ideally, include and relationships wouldn't be parameters of this method, they should only
+        # be used to build the request for the server, and this method would build node according to the response.
+
         nodes: list[InfrahubNode] = []
         related_nodes: list[InfrahubNode] = []
 
@@ -571,7 +574,6 @@ async def _process_nodes_and_relationships(
                     branch=branch,
                     related_nodes=related_nodes,
                     timeout=timeout,
-                    include=include,
                 )
 
         return ProcessRelationsNode(nodes=nodes, related_nodes=related_nodes)
@@ -1837,6 +1839,7 @@ def _process_nodes_and_relationships(
         schema_kind: str,
         branch: str,
         prefetch_relationships: bool,
+        include: list[str] | None,
         timeout: int | None = None,
     ) -> ProcessRelationsNodeSync:
         """Processes InfrahubNodeSync and their Relationships from the GraphQL query response.
@@ -1861,7 +1864,7 @@ def _process_nodes_and_relationships(
             node = InfrahubNodeSync.from_graphql(client=self, branch=branch, data=item, timeout=timeout)
             nodes.append(node)
 
-            if prefetch_relationships:
+            if prefetch_relationships or include is not None:
                 node._process_relationships(node_data=item, branch=branch, related_nodes=related_nodes, timeout=timeout)
 
         return ProcessRelationsNodeSync(nodes=nodes, related_nodes=related_nodes)
@@ -1986,6 +1989,7 @@ def process_page(page_offset: int, page_number: int) -> tuple[dict, ProcessRelat
                 branch=branch,
                 prefetch_relationships=prefetch_relationships,
                 timeout=timeout,
+                include=include,
             )
             return response, process_result
 

infrahub_sdk/node/node.py

@@ -890,7 +890,6 @@ async def _process_relationships(
         branch: str,
         related_nodes: list[InfrahubNode],
         timeout: int | None = None,
-        include: list[str] | None = None,
     ) -> None:
         """Processes the Relationships of a InfrahubNode and add Related Nodes to a list.
 
@@ -901,8 +900,6 @@ async def _process_relationships(
             timeout (int, optional): Overrides default timeout used when querying the graphql API. Specified in seconds.
         """
         for rel_name in self._relationships:
-            if include is not None and rel_name not in include:
-                continue
             rel = getattr(self, rel_name)
             if rel and isinstance(rel, RelatedNode):
                 relation = node_data["node"].get(rel_name, None)
@@ -1369,7 +1366,8 @@ def generate_query_data_node(
                 continue
 
             peer_data: dict[str, Any] = {}
-            if rel_schema and prefetch_relationships:
+            should_fetch_relationship = prefetch_relationships or (include is not None and rel_name in include)
+            if rel_schema and should_fetch_relationship:
                 peer_schema = self._client.schema.get(kind=rel_schema.peer, branch=self._branch)
                 peer_node = InfrahubNodeSync(client=self._client, schema=peer_schema, branch=self._branch)
                 peer_data = peer_node.generate_query_data_node(include=include, exclude=exclude, property=property)

tests/unit/sdk/conftest.py

@@ -2641,3 +2641,21 @@ async def mock_query_tasks_05(httpx_mock: HTTPXMock) -> HTTPXMock:
         is_reusable=True,
     )
     return httpx_mock
+
+
+@pytest.fixture
+async def set_builtin_tag_schema_cache(client) -> None:
+    # Set tag schema in cache to avoid needed to request the server.
+    builtin_tag_schema = {
+        "version": "1.0",
+        "nodes": [
+            {
+                "name": "Tag",
+                "namespace": "Builtin",
+                "default_filter": "name__value",
+                "display_label": "name__value",
+                "branch": "aware",
+            }
+        ],
+    }
+    client.schema.set_cache(builtin_tag_schema)

tests/unit/sdk/test_node.py

@@ -1055,7 +1055,9 @@ async def test_query_data_generic_fragment(clients, mock_schema_query_02, client
 
 
 @pytest.mark.parametrize("client_type", client_types)
-async def test_query_data_include_property(client, location_schema: NodeSchemaAPI, client_type) -> None:
+async def test_query_data_include_property(
+    client, location_schema: NodeSchemaAPI, client_type, set_builtin_tag_schema_cache
+) -> None:
     if client_type == "standard":
         node = InfrahubNode(client=client, schema=location_schema)
         data = await node.generate_query_data(include=["tags"], property=True)
@@ -1178,7 +1180,9 @@ async def test_query_data_include_property(client, location_schema: NodeSchemaAP
 
 
 @pytest.mark.parametrize("client_type", client_types)
-async def test_query_data_include(client, location_schema: NodeSchemaAPI, client_type) -> None:
+async def test_query_data_include(
+    client, location_schema: NodeSchemaAPI, client_type, set_builtin_tag_schema_cache
+) -> None:
     if client_type == "standard":
         node = InfrahubNode(client=client, schema=location_schema)
         data = await node.generate_query_data(include=["tags"])