Repository: mirumee/ariadne-codegen
Pull request #93 (2023-03-09 16:47:40)

More hooks
This pr adds hooks related to:
- generation of enums classes and file
- adding imports to `__init__.py`
- generation of client module, `gql `function, client class, client's methods, and methods arguments


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index cc78c4d..daa6d9c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 - Fixed generating models from interfaces with inline fragments.
 - Added default `None` values for generated methods optional arguments.
 - Added basic plugin system.
+- Added `InitFileGenerator`, `EnumsGenerator`, `ClientGenerator` and `ArgumentsGenerator` plugin hooks.
 
 
 ## 0.3.0 (2023-02-21)
diff --git a/PLUGINS.md b/PLUGINS.md
index 2e9d00a..3e9c300 100644
--- a/PLUGINS.md
+++ b/PLUGINS.md
@@ -31,6 +31,97 @@ def generate_init_module(self, module: ast.Module) -> ast.Module:
 
 Hook executed on generation of init module. Module has list of public, generated classes and reimports them all. Later this module will be saved as `__init__.py`.
 
+### generate_init_import
+
+```py
+def generate_init_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+```
+
+Hook executed on addition of import to init module. Later this import will be placed in `__init__.py`.
+
+### generate_enum
+
+```py
+def generate_enum(self, class_def: ast.ClassDef, enum_type: GraphQLEnumType) -> ast.ClassDef:
+```
+
+Hook executed on generation of class definition of single enum.
+
+### generate_enums_module
+
+```py
+def generate_enums_module(self, module: ast.Module) -> ast.Module:
+```
+
+Hook executed on generation of enums module. Module has all classes representing enums from schema. Later this module will be saved as `{enums_module_name}.py`, `enums_module_name` is taken from config.
+
+### generate_client_module
+
+```py
+def generate_client_module(self, module: ast.Module) -> ast.Module:
+```
+
+Hook executed on generation of client module. Module contains `gql` function definition and client class. Later this module will be saved as `{client_file_name}.py`, `client_file_name` is taken from config.
+
+### generate_gql_function
+
+```py
+def generate_gql_function(self, function_def: ast.FunctionDef) -> ast.FunctionDef:
+```
+
+Hook executed on generation of `gql` function. 
+
+### generate_client_class
+
+```py
+def generate_client_class(self, class_def: ast.ClassDef) -> ast.ClassDef:
+```
+
+Hook executed on generation of client class. Class contains method for every graphql operation.
+
+### generate_client_import
+
+```py
+def generate_client_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+```
+
+Hook executed on addition of import to client module. Later this import will be placed in `{client_file_name}.py`, `client_file_name` is taken from config.
+
+### generate_client_method
+
+```py
+def generate_client_method(
+        self, method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+) -> Union[ast.FunctionDef, ast.AsyncFunctionDef]:
+```
+
+Hook executed on generation of client's method, which represents single graphql operation. Depends on the configuration method can be either async or not.
+
+
+### generate_arguments
+
+```py
+def generate_arguments(
+    self,
+    arguments: ast.arguments,
+    variable_definitions: Tuple[VariableDefinitionNode, ...],
+) -> ast.arguments:
+```
+
+Hook executed on generation of arguments for specific client's method.
+
+### generate_arguments_dict
+
+```py
+def generate_arguments_dict(
+    self,
+    dict_: ast.Dict,
+    variable_definitions: Tuple[VariableDefinitionNode, ...],
+) -> ast.Dict:
+```
+
+Hook executed on generation of dictionary with arguments of graphql operation. Serialized dictionary is later used as variables payload. 
+
 
 ## Example
 
diff --git a/ariadne_codegen/generators/arguments.py b/ariadne_codegen/generators/arguments.py
index 24f5f73..227b6d9 100644
--- a/ariadne_codegen/generators/arguments.py
+++ b/ariadne_codegen/generators/arguments.py
@@ -14,6 +14,7 @@ from graphql import (
 )
 
 from ..exceptions import ParsingError
+from ..plugins.manager import PluginManager
 from .codegen import (
     generate_annotation_name,
     generate_arg,
@@ -35,10 +36,13 @@ class ArgumentsGenerator:
         schema: GraphQLSchema,
         convert_to_snake_case: bool = True,
         custom_scalars: Optional[Dict[str, ScalarData]] = None,
+        plugin_manager: Optional[PluginManager] = None,
     ) -> None:
         self.schema = schema
         self.convert_to_snake_case = convert_to_snake_case
         self.custom_scalars = custom_scalars if custom_scalars else {}
+        self.plugin_manager = plugin_manager
+
         self.used_types: List[str] = []
         self._used_enums: List[str] = []
         self._used_inputs: List[str] = []
@@ -71,6 +75,14 @@ class ArgumentsGenerator:
             args=required_args + optional_args,
             defaults=[generate_constant(None) for _ in optional_args],
         )
+
+        if self.plugin_manager:
+            arguments = self.plugin_manager.generate_arguments(
+                arguments, variable_definitions=variable_definitions
+            )
+            dict_ = self.plugin_manager.generate_arguments_dict(
+                dict_, variable_definitions=variable_definitions
+            )
         return arguments, dict_
 
     def get_used_enums(self) -> List[str]:
diff --git a/ariadne_codegen/generators/client.py b/ariadne_codegen/generators/client.py
index 1073c0b..3077f62 100644
--- a/ariadne_codegen/generators/client.py
+++ b/ariadne_codegen/generators/client.py
@@ -1,6 +1,7 @@
 import ast
-from typing import List
+from typing import List, Optional, Union, cast
 
+from ..plugins.manager import PluginManager
 from .codegen import (
     generate_ann_assign,
     generate_arg,
@@ -15,6 +16,7 @@ from .codegen import (
     generate_import_from,
     generate_keyword,
     generate_method_definition,
+    generate_module,
     generate_name,
     generate_return,
     generate_subscript,
@@ -24,8 +26,14 @@ from .constants import ANY, LIST, OPTIONAL, TYPING_MODULE
 
 
 class ClientGenerator:
-    def __init__(self, name: str, base_client: str) -> None:
+    def __init__(
+        self,
+        name: str,
+        base_client: str,
+        plugin_manager: Optional[PluginManager] = None,
+    ) -> None:
         self.name = name
+        self.plugin_manager = plugin_manager
         self.class_def = generate_class_def(name=name, base_names=[base_client])
         self.imports: list = [
             generate_import_from([OPTIONAL, LIST, ANY], TYPING_MODULE)
@@ -38,20 +46,31 @@ class ClientGenerator:
         self._data_variable = "data"
 
     def generate(self) -> ast.Module:
-        """Generate module with class definition of grahql client."""
+        """Generate module with class definition of graphql client."""
         gql_func = self._generate_gql_func()
         gql_func.lineno = len(self.imports) + 1
+        if self.plugin_manager:
+            gql_func = self.plugin_manager.generate_gql_function(gql_func)
+
         self.class_def.lineno = len(self.imports) + 3
         if not self.class_def.body:
             self.class_def.body.append(ast.Pass())
-        return ast.Module(
+        if self.plugin_manager:
+            self.class_def = self.plugin_manager.generate_client_class(self.class_def)
+
+        module = generate_module(
             body=self.imports + [gql_func, self.class_def],
-            type_ignores=[],
         )
+        if self.plugin_manager:
+            module = self.plugin_manager.generate_client_module(module)
+        return module
 
     def add_import(self, names: List[str], from_: str, level: int = 0) -> None:
         """Add import to be included in module file."""
-        self.imports.append(generate_import_from(names=names, from_=from_, level=level))
+        import_ = generate_import_from(names=names, from_=from_, level=level)
+        if self.plugin_manager:
+            import_ = self.plugin_manager.generate_client_import(import_)
+        self.imports.append(import_)
 
     def add_method(
         self,
@@ -81,6 +100,10 @@ class ClientGenerator:
             )
         )
         method_def.lineno = len(self.class_def.body) + 1
+        if self.plugin_manager:
+            method_def = self.plugin_manager.generate_client_method(
+                cast(Union[ast.FunctionDef, ast.AsyncFunctionDef], method_def)
+            )
         self.class_def.body.append(method_def)
 
     def _generate_async_method(
diff --git a/ariadne_codegen/generators/enums.py b/ariadne_codegen/generators/enums.py
index 2619915..db9db7f 100644
--- a/ariadne_codegen/generators/enums.py
+++ b/ariadne_codegen/generators/enums.py
@@ -1,8 +1,9 @@
 import ast
-from typing import List, cast
+from typing import List, Optional, cast
 
 from graphql import GraphQLEnumType, GraphQLSchema
 
+from ..plugins.manager import PluginManager
 from .codegen import (
     generate_assign,
     generate_class_def,
@@ -14,8 +15,11 @@ from .constants import ENUM_CLASS, ENUM_MODULE
 
 
 class EnumsGenerator:
-    def __init__(self, schema: GraphQLSchema) -> None:
+    def __init__(
+        self, schema: GraphQLSchema, plugin_manager: Optional[PluginManager] = None
+    ) -> None:
         self.schema = schema
+        self.plugin_manager = plugin_manager
 
         self._imports: List[ast.ImportFrom] = [
             generate_import_from([ENUM_CLASS], ENUM_MODULE)
@@ -25,11 +29,13 @@ class EnumsGenerator:
         ]
 
     def generate(self) -> ast.Module:
-        module_body = cast(List[ast.stmt], self._imports) + cast(
-            List[ast.stmt], self._class_defs
+        module = generate_module(
+            body=cast(List[ast.stmt], self._imports)
+            + cast(List[ast.stmt], self._class_defs)
         )
-
-        return generate_module(body=module_body)
+        if self.plugin_manager:
+            module = self.plugin_manager.generate_enums_module(module)
+        return module
 
     def get_generated_public_names(self) -> List[str]:
         return [c.name for c in self._class_defs]
@@ -49,8 +55,11 @@ class EnumsGenerator:
             )
         ]
 
-        return generate_class_def(
+        class_def = generate_class_def(
             name=definition.name,
             base_names=["str", ENUM_CLASS],
             body=cast(List[ast.stmt], fields),
         )
+        if self.plugin_manager:
+            class_def = self.plugin_manager.generate_enum(class_def, definition)
+        return class_def
diff --git a/ariadne_codegen/generators/init_file.py b/ariadne_codegen/generators/init_file.py
index c5cd4dd..257a379 100644
--- a/ariadne_codegen/generators/init_file.py
+++ b/ariadne_codegen/generators/init_file.py
@@ -12,7 +12,10 @@ class InitFileGenerator:
 
     def add_import(self, names: List[str], from_: str, level: int = 0) -> None:
         """Add import to be included in init file."""
-        self.imports.append(generate_import_from(names=names, from_=from_, level=level))
+        import_ = generate_import_from(names=names, from_=from_, level=level)
+        if self.plugin_manager:
+            import_ = self.plugin_manager.generate_init_import(import_)
+        self.imports.append(import_)
 
     def generate(self) -> ast.Module:
         """Generate init with imports and public api of package."""
diff --git a/ariadne_codegen/generators/package.py b/ariadne_codegen/generators/package.py
index 7cc727c..822f368 100644
--- a/ariadne_codegen/generators/package.py
+++ b/ariadne_codegen/generators/package.py
@@ -97,7 +97,9 @@ class PackageGenerator:
             client_generator
             if client_generator
             else ClientGenerator(
-                name=self.client_name, base_client=self.base_client_name
+                name=self.client_name,
+                base_client=self.base_client_name,
+                plugin_manager=self.plugin_manager,
             )
         )
         self.arguments_generator = (
@@ -107,6 +109,7 @@ class PackageGenerator:
                 schema=self.schema,
                 convert_to_snake_case=self.convert_to_snake_case,
                 custom_scalars=self.custom_scalars,
+                plugin_manager=self.plugin_manager,
             )
         )
         self.input_types_generator = (
@@ -121,7 +124,9 @@ class PackageGenerator:
             )
         )
         self.enums_generator = (
-            enums_generator if enums_generator else EnumsGenerator(schema=self.schema)
+            enums_generator
+            if enums_generator
+            else EnumsGenerator(schema=self.schema, plugin_manager=self.plugin_manager)
         )
 
         if base_client_file_path:
diff --git a/ariadne_codegen/plugins/base.py b/ariadne_codegen/plugins/base.py
index eb347a6..0ae4306 100644
--- a/ariadne_codegen/plugins/base.py
+++ b/ariadne_codegen/plugins/base.py
@@ -1,7 +1,7 @@
 import ast
-from typing import Dict
+from typing import Dict, Tuple, Union
 
-from graphql import GraphQLSchema
+from graphql import GraphQLEnumType, GraphQLSchema, VariableDefinitionNode
 
 
 class Plugin:
@@ -11,3 +11,48 @@ class Plugin:
 
     def generate_init_module(self, module: ast.Module) -> ast.Module:
         return module
+
+    def generate_init_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+        return import_
+
+    # pylint: disable=unused-argument
+    def generate_enum(
+        self, class_def: ast.ClassDef, enum_type: GraphQLEnumType
+    ) -> ast.ClassDef:
+        return class_def
+
+    def generate_enums_module(self, module: ast.Module) -> ast.Module:
+        return module
+
+    def generate_client_module(self, module: ast.Module) -> ast.Module:
+        return module
+
+    def generate_gql_function(self, function_def: ast.FunctionDef) -> ast.FunctionDef:
+        return function_def
+
+    def generate_client_class(self, class_def: ast.ClassDef) -> ast.ClassDef:
+        return class_def
+
+    def generate_client_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+        return import_
+
+    def generate_client_method(
+        self, method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+    ) -> Union[ast.FunctionDef, ast.AsyncFunctionDef]:
+        return method_def
+
+    # pylint: disable=unused-argument
+    def generate_arguments(
+        self,
+        arguments: ast.arguments,
+        variable_definitions: Tuple[VariableDefinitionNode, ...],
+    ) -> ast.arguments:
+        return arguments
+
+    # pylint: disable=unused-argument
+    def generate_arguments_dict(
+        self,
+        dict_: ast.Dict,
+        variable_definitions: Tuple[VariableDefinitionNode, ...],
+    ) -> ast.Dict:
+        return dict_
diff --git a/ariadne_codegen/plugins/manager.py b/ariadne_codegen/plugins/manager.py
index f5304b7..3f07fe9 100644
--- a/ariadne_codegen/plugins/manager.py
+++ b/ariadne_codegen/plugins/manager.py
@@ -1,7 +1,7 @@
 import ast
-from typing import Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-from graphql import GraphQLSchema
+from graphql import GraphQLEnumType, GraphQLSchema, VariableDefinitionNode
 
 from .base import Plugin
 
@@ -19,8 +19,62 @@ class PluginManager:
             for cls in plugins_types or []
         ]
 
-    def generate_init_module(self, module: ast.Module) -> ast.Module:
-        modified_module = module
+    def _apply_plugins_on_object(
+        self, method_name: str, obj: Any, *args, **kwargs
+    ) -> Any:
+        modified_obj = obj
         for plugin in self.plugins:
-            modified_module = plugin.generate_init_module(modified_module)
-        return modified_module
+            method = getattr(plugin, method_name)
+            modified_obj = method(modified_obj, *args, **kwargs)
+        return modified_obj
+
+    def generate_init_module(self, module: ast.Module) -> ast.Module:
+        return self._apply_plugins_on_object("generate_init_module", module)
+
+    def generate_init_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+        return self._apply_plugins_on_object("generate_init_import", import_)
+
+    def generate_enum(
+        self, class_def: ast.ClassDef, enum_type: GraphQLEnumType
+    ) -> ast.ClassDef:
+        return self._apply_plugins_on_object(
+            "generate_enum", class_def, enum_type=enum_type
+        )
+
+    def generate_enums_module(self, module: ast.Module) -> ast.Module:
+        return self._apply_plugins_on_object("generate_enums_module", module)
+
+    def generate_client_module(self, module: ast.Module) -> ast.Module:
+        return self._apply_plugins_on_object("generate_client_module", module)
+
+    def generate_gql_function(self, function_def: ast.FunctionDef) -> ast.FunctionDef:
+        return self._apply_plugins_on_object("generate_gql_function", function_def)
+
+    def generate_client_class(self, class_def: ast.ClassDef) -> ast.ClassDef:
+        return self._apply_plugins_on_object("generate_client_class", class_def)
+
+    def generate_client_import(self, import_: ast.ImportFrom) -> ast.ImportFrom:
+        return self._apply_plugins_on_object("generate_client_import", import_)
+
+    def generate_client_method(
+        self, method_def: Union[ast.FunctionDef, ast.AsyncFunctionDef]
+    ) -> Union[ast.FunctionDef, ast.AsyncFunctionDef]:
+        return self._apply_plugins_on_object("generate_client_method", method_def)
+
+    def generate_arguments(
+        self,
+        arguments: ast.arguments,
+        variable_definitions: Tuple[VariableDefinitionNode, ...],
+    ) -> ast.arguments:
+        return self._apply_plugins_on_object(
+            "generate_arguments", arguments, variable_definitions=variable_definitions
+        )
+
+    def generate_arguments_dict(
+        self,
+        dict_: ast.Dict,
+        variable_definitions: Tuple[VariableDefinitionNode, ...],
+    ) -> ast.Dict:
+        return self._apply_plugins_on_object(
+            "generate_arguments_dict", dict_, variable_definitions=variable_definitions
+        )


Test patch:
diff --git a/tests/generators/test_arguments_generator.py b/tests/generators/test_arguments_generator.py
index 622ebc5..d6e0858 100644
--- a/tests/generators/test_arguments_generator.py
+++ b/tests/generators/test_arguments_generator.py
@@ -250,3 +250,37 @@ def test_generate_returns_arguments_with_custom_scalar_and_used_serialize_method
 
     assert compare_ast(arguments, expected_arguments)
     assert compare_ast(arguments_dict, expected_arguments_dict)
+
+
+def test_generate_triggers_generate_arguments_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    schema_str = """
+        schema { query: Query }
+        type Query { _skip: String! }
+        """
+    generator = ArgumentsGenerator(
+        schema=build_schema(schema_str), plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate(
+        _get_variable_definitions_from_query_str("query q($arg: String!) { _skip }")
+    )
+
+    assert mocked_plugin_manager.generate_arguments.called
+
+
+def test_generate_triggers_generate_arguments_dict_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    schema_str = """
+        schema { query: Query }
+        type Query { _skip: String! }
+        """
+    generator = ArgumentsGenerator(
+        schema=build_schema(schema_str), plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate(
+        _get_variable_definitions_from_query_str("query q($arg: String!) { _skip }")
+    )
+
+    assert mocked_plugin_manager.generate_arguments_dict.called
diff --git a/tests/generators/test_client_generator.py b/tests/generators/test_client_generator.py
index 4abf19c..f5f13f8 100644
--- a/tests/generators/test_client_generator.py
+++ b/tests/generators/test_client_generator.py
@@ -41,6 +41,39 @@ def test_generate_returns_module_with_gql_lambda_definition():
     assert compare_ast(assign, expected_def)
 
 
+def test_generate_triggers_generate_gql_function_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = ClientGenerator(
+        "ClientXyz", "BaseClient", plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate()
+
+    assert mocked_plugin_manager.generate_gql_function.called
+
+
+def test_generate_triggers_generate_client_class_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = ClientGenerator(
+        "ClientXyz", "BaseClient", plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate()
+
+    assert mocked_plugin_manager.generate_client_class.called
+
+
+def test_generate_triggers_generate_client_module_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = ClientGenerator(
+        "ClientXyz", "BaseClient", plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate()
+
+    assert mocked_plugin_manager.generate_client_module.called
+
+
 def test_add_import_adds_import_to_generated_module():
     generator = ClientGenerator("Client", "BaseClient")
     name = "Xyz"
@@ -56,6 +89,17 @@ def test_add_import_adds_import_to_generated_module():
     assert [n.name for n in generated_import.names] == [name]
 
 
+def test_add_import_triggers_generate_client_import_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = ClientGenerator(
+        "ClientXyz", "BaseClient", plugin_manager=mocked_plugin_manager
+    )
+
+    generator.add_import(["TestType"], "test", level=1)
+
+    assert mocked_plugin_manager.generate_client_import.called
+
+
 def test_add_method_adds_async_method_definition():
     generator = ClientGenerator("Client", "AsyncBaseClient")
     method_name = "list_xyz"
@@ -304,3 +348,28 @@ def test_add_method_generates_correct_method_body():
     method_def = class_def.body[0]
     assert isinstance(method_def, ast.FunctionDef)
     assert compare_ast(method_def.body, expected_method_body)
+
+
+def test_add_method_triggers_generate_client_method_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = ClientGenerator(
+        "ClientXyz", "BaseClient", plugin_manager=mocked_plugin_manager
+    )
+
+    generator.add_method(
+        name="list_xyz",
+        return_type="ListXyz",
+        arguments=ast.arguments(
+            posonlyargs=[],
+            args=[ast.arg(arg="self")],
+            kwonlyargs=[],
+            kw_defaults=[],
+            defaults=[],
+        ),
+        arguments_dict=ast.Dict(keys=[], values=[]),
+        operation_str="",
+        async_=False,
+    )
+    generator.add_import(["TestType"], "test", level=1)
+
+    assert mocked_plugin_manager.generate_client_method.called
diff --git a/tests/generators/test_enums_generator.py b/tests/generators/test_enums_generator.py
index a2e244f..c62f999 100644
--- a/tests/generators/test_enums_generator.py
+++ b/tests/generators/test_enums_generator.py
@@ -1,6 +1,6 @@
 import ast
 
-from graphql import build_ast_schema, parse
+from graphql import GraphQLEnumType, GraphQLSchema, build_ast_schema, parse
 
 from ariadne_codegen.generators.constants import ENUM_CLASS, ENUM_MODULE
 from ariadne_codegen.generators.enums import EnumsGenerator
@@ -129,3 +129,41 @@ def test_generate_returns_module_with_enum_class_definition_for_every_enum():
     ]
     class_defs = filter_class_defs(module)
     assert compare_ast(class_defs, expected_class_defs)
+
+
+def test_generate_triggers_generate_enums_module_hook(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = EnumsGenerator(
+        schema=GraphQLSchema(), plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate()
+
+    assert mocked_plugin_manager.generate_enums_module.called
+
+
+def test_generate_triggers_generate_enum_hook_for_every_definition(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    schema_str = """
+    enum TestEnumAB {
+        A
+        B
+    }
+    enum TestEnumCD {
+        C
+        D
+    }
+    """
+    generator = EnumsGenerator(
+        schema=build_ast_schema(parse(schema_str)), plugin_manager=mocked_plugin_manager
+    )
+
+    generator.generate()
+
+    assert mocked_plugin_manager.generate_enum.call_count == 2
+    _, call0_enum_type = mocked_plugin_manager.generate_enum.mock_calls[0].args
+    _, call1_enum_type = mocked_plugin_manager.generate_enum.mock_calls[1].args
+    assert isinstance(call0_enum_type, GraphQLEnumType)
+    assert call0_enum_type.name == "TestEnumAB"
+    assert isinstance(call1_enum_type, GraphQLEnumType)
+    assert call1_enum_type.name == "TestEnumCD"
diff --git a/tests/generators/test_init_file_generator.py b/tests/generators/test_init_file_generator.py
index 45f2023..c086343 100644
--- a/tests/generators/test_init_file_generator.py
+++ b/tests/generators/test_init_file_generator.py
@@ -18,6 +18,15 @@ def test_add_import_adds_correct_objects_to_list():
     assert import_.level == 1
 
 
+def test_add_import_triggers_generate_init_import_hook_method(mocker):
+    mocked_plugin_manager = mocker.MagicMock()
+    generator = InitFileGenerator(plugin_manager=mocked_plugin_manager)
+
+    generator.add_import(names=["TestClass"], from_="test", level=1)
+
+    assert mocked_plugin_manager.generate_init_import.called
+
+
 def test_generate_without_imports_returns_empty_module():
     generator = InitFileGenerator()
 
diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py
index 51d644b..ac1f902 100644
--- a/tests/plugins/test_manager.py
+++ b/tests/plugins/test_manager.py
@@ -1,11 +1,26 @@
 import ast
+from typing import cast
 
-from graphql import GraphQLSchema
+import pytest
+from graphql import (
+    GraphQLEnumType,
+    GraphQLEnumValueMap,
+    GraphQLSchema,
+    VariableDefinitionNode,
+)
 
 from ariadne_codegen.plugins.base import Plugin
 from ariadne_codegen.plugins.manager import PluginManager
 
 
+@pytest.fixture
+def mocked_plugin_manager(mocker):
+    return PluginManager(
+        schema=GraphQLSchema(),
+        plugins_types=[mocker.MagicMock(), mocker.MagicMock()],
+    )
+
+
 def test_init_creates_plugins_instances():
     class TestPlugin(Plugin):
         pass
@@ -16,15 +31,101 @@ def test_init_creates_plugins_instances():
     assert isinstance(manager.plugins[0], TestPlugin)
 
 
-def test_generate_init_module_calls_all_plugins(mocker):
-    mocked_plugin_class1 = mocker.MagicMock()
-    mocked_plugin_class2 = mocker.MagicMock()
-    manager = PluginManager(
-        schema=GraphQLSchema(),
-        plugins_types=[mocked_plugin_class1, mocked_plugin_class2],
+def test_generate_init_module_calls_plugins_generate_init_module(mocked_plugin_manager):
+    mocked_plugin_manager.generate_init_module(ast.Module(body=[]))
+
+    assert mocked_plugin_manager.plugins[0].generate_init_module.called
+    assert mocked_plugin_manager.plugins[1].generate_init_module.called
+
+
+def test_generate_init_import_calls_plugins_generate_init_import(mocked_plugin_manager):
+    mocked_plugin_manager.generate_init_import(ast.ImportFrom())
+
+    assert mocked_plugin_manager.plugins[0].generate_init_import.called
+    assert mocked_plugin_manager.plugins[1].generate_init_import.called
+
+
+def test_generate_enum_calls_plugins_generate_enum(mocked_plugin_manager):
+    mocked_plugin_manager.generate_enum(
+        ast.ClassDef(),
+        enum_type=GraphQLEnumType("TestEnum", values=cast(GraphQLEnumValueMap, {})),
+    )
+
+    assert mocked_plugin_manager.plugins[0].generate_enum.called
+    assert mocked_plugin_manager.plugins[1].generate_enum.called
+
+
+def test_generate_enums_module_calls_plugins_generate_enums_module(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_enums_module(ast.Module())
+
+    assert mocked_plugin_manager.plugins[0].generate_enums_module.called
+    assert mocked_plugin_manager.plugins[1].generate_enums_module.called
+
+
+def test_generate_client_module_calls_plugins_generate_client_module(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_client_module(ast.Module())
+
+    assert mocked_plugin_manager.plugins[0].generate_client_module.called
+    assert mocked_plugin_manager.plugins[1].generate_client_module.called
+
+
+def test_generate_gql_function_calls_plugins_generate_gql_function(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_gql_function(ast.FunctionDef())
+
+    assert mocked_plugin_manager.plugins[0].generate_gql_function.called
+    assert mocked_plugin_manager.plugins[1].generate_gql_function.called
+
+
+def test_generate_client_class_calls_plugins_generate_client_class(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_client_class(ast.ClassDef())
+
+    assert mocked_plugin_manager.plugins[0].generate_client_class.called
+    assert mocked_plugin_manager.plugins[1].generate_client_class.called
+
+
+def test_generate_client_import_calls_plugins_generate_client_import(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_client_import(ast.ImportFrom())
+
+    assert mocked_plugin_manager.plugins[0].generate_client_import.called
+    assert mocked_plugin_manager.plugins[1].generate_client_import.called
+
+
+def test_generate_client_method_calls_plugins_generate_client_method(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_client_method(ast.AsyncFunctionDef())
+
+    assert mocked_plugin_manager.plugins[0].generate_client_method.called
+    assert mocked_plugin_manager.plugins[1].generate_client_method.called
+
+
+def test_generate_arguments_calls_plugins_generate_arguments(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_arguments(
+        ast.arguments(), variable_definitions=(VariableDefinitionNode(),)
     )
 
-    manager.generate_init_module(ast.Module(body=[]))
+    assert mocked_plugin_manager.plugins[0].generate_arguments.called
+    assert mocked_plugin_manager.plugins[1].generate_arguments.called
+
+
+def test_generate_arguments_dict_calls_plugins_generate_arguments_dict(
+    mocked_plugin_manager,
+):
+    mocked_plugin_manager.generate_arguments_dict(
+        ast.Dict(), variable_definitions=(VariableDefinitionNode(),)
+    )
 
-    assert manager.plugins[0].generate_init_module.called
-    assert manager.plugins[1].generate_init_module.called
+    assert mocked_plugin_manager.plugins[0].generate_arguments_dict.called
+    assert mocked_plugin_manager.plugins[1].generate_arguments_dict.called