Add unit test, breakdown plugin into multiple plugins

Author: dgarros

infrahub_sdk/ctl/graphql.py

@@ -9,6 +9,7 @@
 from ariadne_codegen.plugins.manager import PluginManager
 from ariadne_codegen.schema import (
     filter_operations_definitions,
+    filter_fragments_definitions,
     get_graphql_schema_from_path,
 )
 from ariadne_codegen.settings import ClientSettings, CommentsStrategy
@@ -23,7 +24,11 @@
 app = AsyncTyper()
 console = Console()
 
-ARIADNE_PLUGINS = ["infrahub_sdk.graphql.plugin.InfrahubPlugin"]
+ARIADNE_PLUGINS = [
+    "infrahub_sdk.graphql.plugin.PydanticBaseModelPlugin",
+    "infrahub_sdk.graphql.plugin.FutureAnnotationPlugin",
+    "infrahub_sdk.graphql.plugin.StandardTypeHintPlugin",
+]
 
 
 def find_gql_files(query_path: Path) -> list[Path]:
@@ -108,33 +113,36 @@ async def generate_return_types(
     for gql_file in gql_files:
         gql_per_directory[gql_file.parent].append(gql_file)
 
+
     # Generate the Pydantic Models for the GraphQL queries
     for directory, gql_files in gql_per_directory.items():
-        package_generator = get_package_generator(
-            schema=graphql_schema,
-            fragments=[],
-            settings=ClientSettings(
-                schema_path=str(schema),
-                target_package_name=directory.name,
-                queries_path=str(directory),
-                include_comments=CommentsStrategy.NONE,
-            ),
-            plugin_manager=plugin_manager,
-        )
-
         for gql_file in gql_files:
             try:
                 definitions = get_graphql_query(queries_path=gql_file, schema=graphql_schema)
             except ValueError as e:
                 print(f"Error generating result types for {gql_file}: {e}")
                 continue
             queries = filter_operations_definitions(definitions)
+            fragments = filter_fragments_definitions(definitions)
+
+            package_generator = get_package_generator(
+                schema=graphql_schema,
+                fragments=fragments,
+                settings=ClientSettings(
+                    schema_path=str(schema),
+                    target_package_name=directory.name,
+                    queries_path=str(directory),
+                    include_comments=CommentsStrategy.NONE,
+                ),
+                plugin_manager=plugin_manager,
+            )
+
 
             for query_operation in queries:
                 package_generator.add_operation(query_operation)
 
-        # package_generator._generate_result_types()
-        generate_result_types(directory=directory, package=package_generator)
+            package_generator._generate_fragments()
+            generate_result_types(directory=directory, package=package_generator)
 
-        for file_name in package_generator._result_types_files.keys():
-            print(f"Generated {file_name} in {directory}")
+            for file_name in package_generator._result_types_files.keys():
+                print(f"Generated {file_name} in {directory}")

infrahub_sdk/graphql/plugin.py

@@ -5,48 +5,32 @@
 
 from ariadne_codegen.plugins.base import Plugin
 
-from .return_type import GraphQLReturnTypeModel
-
 if TYPE_CHECKING:
     from graphql import ExecutableDefinitionNode
 
 
-class InfrahubPlugin(Plugin):
+class FutureAnnotationPlugin(Plugin):
     @staticmethod
-    def find_base_model_index(module: ast.Module) -> int:
-        for idx, item in enumerate(module.body):
-            if isinstance(item, ast.ImportFrom) and item.module == "base_model":
-                return idx
-        return -1
+    def insert_future_annotation(module: ast.Module) -> ast.Module:
+        # First check if the future annotation is already present
+        for item in module.body:
+            if isinstance(item, ast.ImportFrom) and item.module == "__future__" and item.names[0].name == "annotations":
+                return module
 
-    @classmethod
-    def replace_base_model_import(cls, module: ast.Module) -> ast.Module:
-        base_model_index = cls.find_base_model_index(module)
-        if base_model_index == -1:
-            raise ValueError("BaseModel not found in module")
-        module.body[base_model_index] = ast.ImportFrom(
-            module="infrahub_sdk.graphql", names=[ast.alias(name=GraphQLReturnTypeModel.__name__)], level=2
-        )
+        module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0))
         return module
 
-    @staticmethod
-    def replace_base_model_class(module: ast.Module) -> ast.Module:
-        """Replace the BaseModel inserted by Ariadne with the GraphQLReturnTypeModel class."""
-        for item in module.body:
-            if not isinstance(item, ast.ClassDef):
-                continue
+    def generate_result_types_module(
+        self,
+        module: ast.Module,
+        operation_definition: ExecutableDefinitionNode,  # noqa: ARG002
+    ) -> ast.Module:
+        module = self.insert_future_annotation(module)
 
-            for base in item.bases:
-                if isinstance(base, ast.Name) and base.id == "BaseModel":
-                    base.id = GraphQLReturnTypeModel.__name__
         return module
 
-    @staticmethod
-    def insert_future_annotation(module: ast.Module) -> ast.Module:
-        """Insert the future annotation at the beginning of the module."""
-        module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0))
-        return module
 
+class StandardTypeHintPlugin(Plugin):
     @classmethod
     def replace_list_in_subscript(cls, subscript: ast.Subscript) -> ast.Subscript:
         if isinstance(subscript.value, ast.Name) and subscript.value.id == "List":
@@ -76,10 +60,33 @@ def generate_result_types_module(
         module: ast.Module,
         operation_definition: ExecutableDefinitionNode,  # noqa: ARG002
     ) -> ast.Module:
-        module = self.insert_future_annotation(module)
-        module = self.replace_base_model_import(module)
-        module = self.replace_base_model_class(module)
-
+        module = FutureAnnotationPlugin.insert_future_annotation(module)
         module = self.replace_list_annotations(module)
 
         return module
+
+
+class PydanticBaseModelPlugin(Plugin):
+    @staticmethod
+    def find_base_model_index(module: ast.Module) -> int:
+        for idx, item in enumerate(module.body):
+            if isinstance(item, ast.ImportFrom) and item.module == "base_model":
+                return idx
+        return -1
+
+    @classmethod
+    def replace_base_model_import(cls, module: ast.Module) -> ast.Module:
+        base_model_index = cls.find_base_model_index(module)
+        if base_model_index == -1:
+            raise ValueError("BaseModel not found in module")
+        module.body[base_model_index] = ast.ImportFrom(module="pydantic", names=[ast.alias(name="BaseModel")], level=0)
+        return module
+
+    def generate_result_types_module(
+        self,
+        module: ast.Module,
+        operation_definition: ExecutableDefinitionNode,  # noqa: ARG002
+    ) -> ast.Module:
+        module = self.replace_base_model_import(module)
+
+        return module

pyproject.toml

@@ -132,6 +132,7 @@ exclude = [
     "build",
     "dist",
     "examples",
+    "tests/fixtures/unit/test_graphql_plugin",
 ]
 
 
@@ -257,7 +258,6 @@ max-complexity = 17
     "S105", # 'PASS' is not a password but a state
 ]
 
-
 "tests/**/*.py" = [
     "PLR2004", # Magic value used in comparison
     "S101",    # Use of assert detected
@@ -277,6 +277,11 @@ max-complexity = 17
 "tests/unit/sdk/test_client.py" = [
     "W293", # Blank line contains whitespace (used within output check)
 ]
+
+"tests/fixtures/unit/test_graphql_plugin/*.py" = [
+    "FA100", # Add `from __future__ import annotations` to simplify `typing.Optional`
+]
+
 "tasks.py" = [
     "PLC0415", # `import` should be at the top-level of a file
 ]

tests/fixtures/unit/test_graphql_plugin/python01.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from infrahub_sdk.graphql import GraphQLReturnTypeModel
+
+
+class CreateDevice(GraphQLReturnTypeModel):
+    infra_device_upsert: Optional[CreateDeviceInfraDeviceUpsert] = Field(alias="InfraDeviceUpsert")
+
+
+class CreateDeviceInfraDeviceUpsert(GraphQLReturnTypeModel):
+    ok: Optional[bool]
+    object: Optional[CreateDeviceInfraDeviceUpsertObject]
+
+
+class CreateDeviceInfraDeviceUpsertObject(GraphQLReturnTypeModel):
+    id: str
+    name: Optional[CreateDeviceInfraDeviceUpsertObjectName]
+    description: Optional[CreateDeviceInfraDeviceUpsertObjectDescription]
+    status: Optional[CreateDeviceInfraDeviceUpsertObjectStatus]
+
+
+class CreateDeviceInfraDeviceUpsertObjectName(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+class CreateDeviceInfraDeviceUpsertObjectDescription(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+class CreateDeviceInfraDeviceUpsertObjectStatus(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+CreateDevice.model_rebuild()
+CreateDeviceInfraDeviceUpsert.model_rebuild()
+CreateDeviceInfraDeviceUpsertObject.model_rebuild()

tests/fixtures/unit/test_graphql_plugin/python01_after_annotation.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from infrahub_sdk.graphql import GraphQLReturnTypeModel
+
+
+class CreateDevice(GraphQLReturnTypeModel):
+    infra_device_upsert: Optional[CreateDeviceInfraDeviceUpsert] = Field(alias="InfraDeviceUpsert")
+
+
+class CreateDeviceInfraDeviceUpsert(GraphQLReturnTypeModel):
+    ok: Optional[bool]
+    object: Optional[CreateDeviceInfraDeviceUpsertObject]
+
+
+class CreateDeviceInfraDeviceUpsertObject(GraphQLReturnTypeModel):
+    id: str
+    name: Optional[CreateDeviceInfraDeviceUpsertObjectName]
+    description: Optional[CreateDeviceInfraDeviceUpsertObjectDescription]
+    status: Optional[CreateDeviceInfraDeviceUpsertObjectStatus]
+
+
+class CreateDeviceInfraDeviceUpsertObjectName(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+class CreateDeviceInfraDeviceUpsertObjectDescription(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+class CreateDeviceInfraDeviceUpsertObjectStatus(GraphQLReturnTypeModel):
+    value: Optional[str]
+
+
+CreateDevice.model_rebuild()
+CreateDeviceInfraDeviceUpsert.model_rebuild()
+CreateDeviceInfraDeviceUpsertObject.model_rebuild()

tests/fixtures/unit/test_graphql_plugin/python02.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from pydantic import Field
+
+from infrahub_sdk.graphql import GraphQLReturnTypeModel
+
+
+class CreateDevice(GraphQLReturnTypeModel):
+    infra_device_upsert: Optional["CreateDeviceInfraDeviceUpsert"] = Field(alias="InfraDeviceUpsert")
+
+
+class CreateDeviceInfraDeviceUpsert(GraphQLReturnTypeModel):
+    ok: Optional[bool]
+    object: Optional[dict]
+
+
+CreateDevice.model_rebuild()
+CreateDeviceInfraDeviceUpsert.model_rebuild()

tests/fixtures/unit/test_graphql_plugin/python02_after_annotation.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import Field
+
+from infrahub_sdk.graphql import GraphQLReturnTypeModel
+
+
+class CreateDevice(GraphQLReturnTypeModel):
+    infra_device_upsert: Optional["CreateDeviceInfraDeviceUpsert"] = Field(alias="InfraDeviceUpsert")
+
+
+class CreateDeviceInfraDeviceUpsert(GraphQLReturnTypeModel):
+    ok: Optional[bool]
+    object: Optional[dict]
+
+
+CreateDevice.model_rebuild()
+CreateDeviceInfraDeviceUpsert.model_rebuild()