Adjust set_module_args() for data tagging.

Author: felixfontein

tests/unit/plugins/modules/utils.py

@@ -21,9 +21,17 @@ def set_module_args(args):
     if '_ansible_keep_remote_files' not in args:
         args['_ansible_keep_remote_files'] = False
 
-    serialized_args = to_bytes(json.dumps({'ANSIBLE_MODULE_ARGS': args}))
-    with patch.object(basic, '_ANSIBLE_ARGS', serialized_args):
-        yield
+    try:
+        from ansible.module_utils.testing import patch_module_args
+    except ImportError:
+        # Before data tagging support was merged, this was the way to go:
+        serialized_args = to_bytes(json.dumps({'ANSIBLE_MODULE_ARGS': args}))
+        with patch.object(basic, '_ANSIBLE_ARGS', serialized_args):
+            yield
+    else:
+        # With data tagging support, we have a new helper for this:
+        with patch_module_args(args):
+            yield
 
 
 class AnsibleExitJson(Exception):