From 0cd0f0eaf670d8da8afc6137ee19a53439e06e48 Mon Sep 17 00:00:00 2001
From: Alexei Znamensky <103110+russoz@users.noreply.github.com>
Date: Thu, 8 Apr 2021 17:58:39 +1200
Subject: [PATCH] module_helper - fixed decorator cause_changes (#2203)

* fixed decorator cause_changes

* added changelog fragment

* typo
---
 .../2203-modhelper-cause-changes-deco.yml     |  2 +
 plugins/module_utils/module_helper.py         | 33 ++++++++-----
 .../module_utils/test_module_helper.py        | 46 ++++++++++++++++++-
 3 files changed, 68 insertions(+), 13 deletions(-)
 create mode 100644 changelogs/fragments/2203-modhelper-cause-changes-deco.yml

diff --git a/changelogs/fragments/2203-modhelper-cause-changes-deco.yml b/changelogs/fragments/2203-modhelper-cause-changes-deco.yml
new file mode 100644
index 0000000000..b61f97d6b8
--- /dev/null
+++ b/changelogs/fragments/2203-modhelper-cause-changes-deco.yml
@@ -0,0 +1,2 @@
+bugfixes:
+  - module_helper module utils - fixed decorator ``cause_changes`` (https://github.com/ansible-collections/community.general/pull/2203).
diff --git a/plugins/module_utils/module_helper.py b/plugins/module_utils/module_helper.py
index 6357eae25c..d241eba5af 100644
--- a/plugins/module_utils/module_helper.py
+++ b/plugins/module_utils/module_helper.py
@@ -101,18 +101,27 @@ class ArgFormat(object):
         return [str(p) for p in func(value)]
 
 
-def cause_changes(func, on_success=True, on_failure=False):
-    @wraps(func)
-    def wrapper(self, *args, **kwargs):
-        try:
-            func(*args, **kwargs)
-            if on_success:
-                self.changed = True
-        except Exception:
-            if on_failure:
-                self.changed = True
-            raise
-    return wrapper
+def cause_changes(on_success=None, on_failure=None):
+
+    def deco(func):
+        if on_success is None and on_failure is None:
+            return func
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                self = args[0]
+                func(*args, **kwargs)
+                if on_success is not None:
+                    self.changed = on_success
+            except Exception:
+                if on_failure is not None:
+                    self.changed = on_failure
+                raise
+
+        return wrapper
+
+    return deco
 
 
 def module_fails_on_exception(func):
diff --git a/tests/unit/plugins/module_utils/test_module_helper.py b/tests/unit/plugins/module_utils/test_module_helper.py
index b8ea36501c..6f77ca7662 100644
--- a/tests/unit/plugins/module_utils/test_module_helper.py
+++ b/tests/unit/plugins/module_utils/test_module_helper.py
@@ -6,10 +6,12 @@
 from __future__ import (absolute_import, division, print_function)
 __metaclass__ = type
 
+from collections import namedtuple
+
 import pytest
 
 from ansible_collections.community.general.plugins.module_utils.module_helper import (
-    ArgFormat, DependencyCtxMgr, ModuleHelper, VarMeta
+    ArgFormat, DependencyCtxMgr, ModuleHelper, VarMeta, cause_changes
 )
 
 
@@ -160,3 +162,45 @@ def test_vardict():
     assert vd.c == 'new_c'
     assert vd.output() == {'a': 'new_a', 'c': 'new_c'}
     assert vd.diff() == {'before': {'a': 123}, 'after': {'a': 'new_a'}}, "diff={0}".format(vd.diff())
+
+
+class MockMH(object):
+    changed = None
+
+    def _div(self, x, y):
+        return x / y
+
+    func_none = cause_changes()(_div)
+    func_onsucc = cause_changes(on_success=True)(_div)
+    func_onfail = cause_changes(on_failure=True)(_div)
+    func_onboth = cause_changes(on_success=True, on_failure=True)(_div)
+
+
+CAUSE_CHG_DECO_PARAMS = ['method', 'expect_exception', 'expect_changed']
+CAUSE_CHG_DECO = dict(
+    none_succ=dict(method='func_none', expect_exception=False, expect_changed=None),
+    none_fail=dict(method='func_none', expect_exception=True, expect_changed=None),
+    onsucc_succ=dict(method='func_onsucc', expect_exception=False, expect_changed=True),
+    onsucc_fail=dict(method='func_onsucc', expect_exception=True, expect_changed=None),
+    onfail_succ=dict(method='func_onfail', expect_exception=False, expect_changed=None),
+    onfail_fail=dict(method='func_onfail', expect_exception=True, expect_changed=True),
+    onboth_succ=dict(method='func_onboth', expect_exception=False, expect_changed=True),
+    onboth_fail=dict(method='func_onboth', expect_exception=True, expect_changed=True),
+)
+CAUSE_CHG_DECO_IDS = sorted(CAUSE_CHG_DECO.keys())
+
+
+@pytest.mark.parametrize(CAUSE_CHG_DECO_PARAMS,
+                         [[CAUSE_CHG_DECO[tc][param]
+                           for param in CAUSE_CHG_DECO_PARAMS]
+                          for tc in CAUSE_CHG_DECO_IDS],
+                         ids=CAUSE_CHG_DECO_IDS)
+def test_cause_changes_deco(method, expect_exception, expect_changed):
+    mh = MockMH()
+    if expect_exception:
+        with pytest.raises(Exception):
+            getattr(mh, method)(1, 0)
+    else:
+        getattr(mh, method)(9, 3)
+
+    assert mh.changed == expect_changed