From cfea62793f3008e0540e86828faf41ec1e703aef Mon Sep 17 00:00:00 2001
From: "patchback[bot]" <45432694+patchback[bot]@users.noreply.github.com>
Date: Sun, 5 Dec 2021 22:14:08 +0100
Subject: [PATCH] MH decorators - added decorators for check_mode (#3849)
 (#3860)

* MH decorators - added decorators for check_mode

* added changelog fragment

(cherry picked from commit fb79c2998ea705f0e2f6376706c0cccfb16973d6)

Co-authored-by: Alexei Znamensky <103110+russoz@users.noreply.github.com>
---
 .../fragments/3849-mh-check-mode-decos.yaml   |  2 ++
 plugins/module_utils/mh/deco.py               | 33 +++++++++++++++++++
 .../targets/module_helper/library/msimple.py  | 13 ++++++--
 .../targets/module_helper/tasks/msimple.yml   | 29 +++++++++++++++-
 4 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 changelogs/fragments/3849-mh-check-mode-decos.yaml

diff --git a/changelogs/fragments/3849-mh-check-mode-decos.yaml b/changelogs/fragments/3849-mh-check-mode-decos.yaml
new file mode 100644
index 0000000000..48ba0119a4
--- /dev/null
+++ b/changelogs/fragments/3849-mh-check-mode-decos.yaml
@@ -0,0 +1,2 @@
+minor_changes:
+    - module_helper module utils - added decorators ``check_mode_skip`` and ``check_mode_skip_returns`` for skipping methods when ``check_mode=True`` (https://github.com/ansible-collections/community.general/pull/3849).
diff --git a/plugins/module_utils/mh/deco.py b/plugins/module_utils/mh/deco.py
index 91f0d97744..62d460b4e6 100644
--- a/plugins/module_utils/mh/deco.py
+++ b/plugins/module_utils/mh/deco.py
@@ -52,3 +52,36 @@ def module_fails_on_exception(func):
             self.module.fail_json(msg=msg, exception=traceback.format_exc(),
                                   output=self.output, vars=self.vars.output(), **self.output)
     return wrapper
+
+
+def check_mode_skip(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not self.module.check_mode:
+            return func(self, *args, **kwargs)
+    return wrapper
+
+
+def check_mode_skip_returns(callable=None, value=None):
+
+    def deco(func):
+        if callable is not None:
+            @wraps(func)
+            def wrapper_callable(self, *args, **kwargs):
+                if self.module.check_mode:
+                    return callable(self, *args, **kwargs)
+                return func(self, *args, **kwargs)
+            return wrapper_callable
+
+        if value is not None:
+            @wraps(func)
+            def wrapper_value(self, *args, **kwargs):
+                if self.module.check_mode:
+                    return value
+                return func(self, *args, **kwargs)
+            return wrapper_value
+
+    if callable is None and value is None:
+        return check_mode_skip
+
+    return deco
diff --git a/tests/integration/targets/module_helper/library/msimple.py b/tests/integration/targets/module_helper/library/msimple.py
index da43eca777..bde7f4fcc5 100644
--- a/tests/integration/targets/module_helper/library/msimple.py
+++ b/tests/integration/targets/module_helper/library/msimple.py
@@ -30,30 +30,39 @@ EXAMPLES = ""
 RETURN = ""
 
 from ansible_collections.community.general.plugins.module_utils.module_helper import ModuleHelper
+from ansible_collections.community.general.plugins.module_utils.mh.deco import check_mode_skip
 
 
 class MSimple(ModuleHelper):
     output_params = ('a', 'b', 'c')
     module = dict(
         argument_spec=dict(
-            a=dict(type='int'),
+            a=dict(type='int', default=0),
             b=dict(type='str'),
             c=dict(type='str'),
         ),
+        supports_check_mode=True,
     )
 
     def __init_module__(self):
         self.vars.set('value', None)
         self.vars.set('abc', "abc", diff=True)
 
+    @check_mode_skip
+    def process_a3_bc(self):
+        if self.vars.a == 3:
+            self.vars['b'] = str(self.vars.b) * 3
+            self.vars['c'] = str(self.vars.c) * 3
+
     def __run__(self):
-        if (0 if self.vars.a is None else self.vars.a) >= 100:
+        if self.vars.a >= 100:
             raise Exception("a >= 100")
         if self.vars.c == "abc change":
             self.vars['abc'] = "changed abc"
         if self.vars.get('a', 0) == 2:
             self.vars['b'] = str(self.vars.b) * 2
             self.vars['c'] = str(self.vars.c) * 2
+        self.process_a3_bc()
 
 
 def main():
diff --git a/tests/integration/targets/module_helper/tasks/msimple.yml b/tests/integration/targets/module_helper/tasks/msimple.yml
index 4d2ff9b798..d539fa3f9f 100644
--- a/tests/integration/targets/module_helper/tasks/msimple.yml
+++ b/tests/integration/targets/module_helper/tasks/msimple.yml
@@ -1,7 +1,7 @@
 # (c) 2021, Alexei Znamensky
 # GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
 
-- name: test msimple 1
+- name: test msimple (set a=80)
   msimple:
     a: 80
   register: simple1
@@ -55,3 +55,30 @@
       - simple4.c == "abc change"
       - simple4.abc == "changed abc"
       - simple4 is changed
+
+- name: test msimple 5a
+  msimple:
+    a: 3   # should triple b and c
+    b: oh
+    c: my
+  register: simple5a
+
+- name: test msimple 5b
+  check_mode: true
+  msimple:
+    a: 3   # should triple b and c
+    b: oh
+    c: my
+  register: simple5b
+
+- name: assert simple5
+  assert:
+    that:
+      - simple5a.a == 3
+      - simple5a.b == "ohohoh"
+      - simple5a.c == "mymymy"
+      - simple5a is not changed
+      - simple5b.a == 3
+      - simple5b.b == "oh"
+      - simple5b.c == "my"
+      - simple5b is not changed