Add toggle to fix module_defaults with module-as-redirected-action on a per-module basis (#77265)

* If there is a platform specific handler, prefer the resolved module over the resolved action when loading module_defaults

Add a toggle for action plugins to prefer the resolved module when loading module_defaults

Allow moving away from modules intercepted as actions pattern

Fixes #77059

changelogs/fragments/77265-module_defaults-with-modules-as-redirected-actions.yaml

@@ -0,0 +1,55 @@
+minor_changes:
+- |
+  Add an 'action_plugin' field for modules in runtime.yml plugin_routing.
+
+  This fixes module_defaults by supporting modules-as-redirected-actions
+  without redirecting module_defaults entries to the common action.
+
+  .. code: yaml
+
+     plugin_routing:
+       action:
+         facts:
+           redirect: ns.coll.eos
+         command:
+           redirect: ns.coll.eos
+       modules:
+         facts:
+           redirect: ns.coll.eos_facts
+         command:
+           redirect: ns.coll.eos_command
+
+  With the runtime.yml above for ns.coll, a task such as
+
+  .. code: yaml
+
+     - hosts: all
+       module_defaults:
+         ns.coll.eos_facts: {'valid_for_eos_facts': 'value'}
+         ns.coll.eos_command: {'not_valid_for_eos_facts': 'value'}
+       tasks:
+         - ns.coll.facts:
+
+  will end up with defaults for eos_facts and eos_command
+  since both modules redirect to the same action.
+
+  To select an action plugin for a module without merging
+  module_defaults, define an action_plugin field for the resolved
+  module in the runtime.yml.
+
+  .. code: yaml
+
+     plugin_routing:
+       modules:
+         facts:
+           redirect: ns.coll.eos_facts
+           action_plugin: ns.coll.eos
+         command:
+           redirect: ns.coll.eos_command
+           action_plugin: ns.coll.eos
+
+  The action_plugin field can be a redirected action plugin, as
+  it is resolved normally.
+
+  Using the modified runtime.yml, the example task will only use
+  the ns.coll.eos_facts defaults.

lib/ansible/executor/task_executor.py

@@ -26,7 +26,7 @@ from ansible.playbook.conditional import Conditional
 from ansible.playbook.task import Task
 from ansible.plugins.loader import become_loader, cliconf_loader, connection_loader, httpapi_loader, netconf_loader, terminal_loader
 from ansible.template import Templar
-from ansible.utils.collection_loader import AnsibleCollectionConfig
+from ansible.utils.collection_loader import AnsibleCollectionConfig, AnsibleCollectionRef
 from ansible.utils.listify import listify_lookup_plugin_terms
 from ansible.utils.unsafe_proxy import to_unsafe_text, wrap_var
 from ansible.vars.clean import namespace_facts, clean_facts
@@ -590,11 +590,16 @@ class TaskExecutor:
             cvars['ansible_python_interpreter'] = sys.executable
 
         # get handler
-        self._handler = self._get_action_handler(connection=self._connection, templar=templar)
+        self._handler, module_context = self._get_action_handler_with_module_context(connection=self._connection, templar=templar)
+
+        if module_context is not None:
+            module_defaults_fqcn = module_context.resolved_fqcn
+        else:
+            module_defaults_fqcn = self._task.resolved_action
 
         # Apply default params for action/module, if present
         self._task.args = get_action_args_with_defaults(
-            self._task.resolved_action, self._task.args, self._task.module_defaults, templar,
+            module_defaults_fqcn, self._task.args, self._task.module_defaults, templar,
             action_groups=self._task._parent._play._action_groups
         )
 
@@ -1093,7 +1098,12 @@ class TaskExecutor:
         '''
         Returns the correct action plugin to handle the requestion task action
         '''
+        return self._get_action_handler_with_module_context(connection, templar)[0]
 
+    def _get_action_handler_with_module_context(self, connection, templar):
+        '''
+        Returns the correct action plugin to handle the requestion task action and the module context
+        '''
         module_collection, separator, module_name = self._task.action.rpartition(".")
         module_prefix = module_name.split('_')[0]
         if module_collection:
@@ -1106,8 +1116,16 @@ class TaskExecutor:
 
         collections = self._task.collections
 
+        # Check if the module has specified an action handler
+        module = self._shared_loader_obj.module_loader.find_plugin_with_context(
+            self._task.action, collection_list=collections
+        )
+        if not module.resolved or not module.action_plugin:
+            module = None
+        if module is not None:
+            handler_name = module.action_plugin
         # let action plugin override module, fallback to 'normal' action plugin otherwise
-        if self._shared_loader_obj.action_loader.has_plugin(self._task.action, collection_list=collections):
+        elif self._shared_loader_obj.action_loader.has_plugin(self._task.action, collection_list=collections):
             handler_name = self._task.action
         elif all((module_prefix in C.NETWORK_GROUP_MODULES, self._shared_loader_obj.action_loader.has_plugin(network_action, collection_list=collections))):
             handler_name = network_action
@@ -1133,7 +1151,7 @@ class TaskExecutor:
         if not handler:
             raise AnsibleError("the handler '%s' was not found" % handler_name)
 
-        return handler
+        return handler, module
 
 
 def start_connection(play_context, variables, task_uuid):

lib/ansible/playbook/base.py

@@ -507,9 +507,13 @@ class FieldAttributeBase(metaclass=BaseMeta):
         return fq_group_name, resolved_actions
 
     def _resolve_action(self, action_name, mandatory=True):
-        context = action_loader.find_plugin_with_context(action_name)
-        if not context.resolved:
-            context = module_loader.find_plugin_with_context(action_name)
+        context = module_loader.find_plugin_with_context(action_name)
+        if context.resolved and not context.action_plugin:
+            prefer = action_loader.find_plugin_with_context(action_name)
+            if prefer.resolved:
+                context = prefer
+        elif not context.resolved:
+            context = action_loader.find_plugin_with_context(action_name)
 
         if context.resolved:
             return context.resolved_fqcn

lib/ansible/plugins/loader.py

@@ -126,6 +126,7 @@ class PluginLoadContext(object):
         self.deprecation_warnings = []
         self.resolved = False
         self._resolved_fqcn = None
+        self.action_plugin = None
 
     @property
     def resolved_fqcn(self):
@@ -166,13 +167,14 @@ class PluginLoadContext(object):
         self.deprecation_warnings.append(warning_text)
         return self
 
-    def resolve(self, resolved_name, resolved_path, resolved_collection, exit_reason):
+    def resolve(self, resolved_name, resolved_path, resolved_collection, exit_reason, action_plugin):
         self.pending_redirect = None
         self.plugin_resolved_name = resolved_name
         self.plugin_resolved_path = resolved_path
         self.plugin_resolved_collection = resolved_collection
         self.exit_reason = exit_reason
         self.resolved = True
+        self.action_plugin = action_plugin
         return self
 
     def redirect(self, redirect_name):
@@ -231,8 +233,12 @@ class PluginLoader:
 
         self._searched_paths = set()
 
+    @property
+    def type(self):
+        return AnsibleCollectionRef.legacy_plugin_dir_to_plugin_type(self.subdir)
+
     def __repr__(self):
-        return 'PluginLoader(type={0})'.format(AnsibleCollectionRef.legacy_plugin_dir_to_plugin_type(self.subdir))
+        return 'PluginLoader(type={0})'.format(self.type)
 
     def _clear_caches(self):
 
@@ -459,6 +465,7 @@ class PluginLoader:
         # check collection metadata to see if any special handling is required for this plugin
         routing_metadata = self._query_collection_routing_meta(acr, plugin_type, extension=extension)
 
+        action_plugin = None
         # TODO: factor this into a wrapper method
         if routing_metadata:
             deprecation = routing_metadata.get('deprecation', None)
@@ -497,6 +504,9 @@ class PluginLoader:
                 return plugin_load_context.redirect(redirect)
                 # TODO: non-FQCN case, do we support `.` prefix for current collection, assume it with no dots, require it for subdirs in current, or ?
 
+            if self.type == 'modules':
+                action_plugin = routing_metadata.get('action_plugin')
+
         n_resource = to_native(acr.resource, errors='strict')
         # we want this before the extension is added
         full_name = '{0}.{1}'.format(acr.n_python_package_name, n_resource)
@@ -519,7 +529,7 @@ class PluginLoader:
         # FIXME: and is file or file link or ...
         if os.path.exists(n_resource_path):
             return plugin_load_context.resolve(
-                full_name, to_text(n_resource_path), acr.collection, 'found exact match for {0} in {1}'.format(full_name, acr.collection))
+                full_name, to_text(n_resource_path), acr.collection, 'found exact match for {0} in {1}'.format(full_name, acr.collection), action_plugin)
 
         if extension:
             # the request was extension-specific, don't try for an extensionless match
@@ -538,7 +548,8 @@ class PluginLoader:
             pass
 
         return plugin_load_context.resolve(
-            full_name, to_text(found_files[0]), acr.collection, 'found fuzzy extension match for {0} in {1}'.format(full_name, acr.collection))
+            full_name, to_text(found_files[0]), acr.collection,
+            'found fuzzy extension match for {0} in {1}'.format(full_name, acr.collection), action_plugin)
 
     def find_plugin(self, name, mod_type='', ignore_deprecated=False, check_aliases=False, collection_list=None):
         ''' Find a plugin named name '''

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/meta/runtime.yml

@@ -1,3 +1,49 @@
+plugin_routing:
+  action:
+    # Backwards compat for modules-redirected-as-actions:
+    # By default, each module_defaults entry is resolved as an action plugin,
+    # and if it does not exist, it is resolved a a module.
+    # All modules that redirect to the same action will resolve to the same action.
+    module_uses_action_defaults:
+      redirect: testns.testcoll.eos
+
+    # module-redirected-as-action overridden by action_plugin
+    iosfacts:
+      redirect: testns.testcoll.nope
+    ios_facts:
+      redirect: testns.testcoll.nope
+
+    redirected_action:
+      redirect: testns.testcoll.ios
+  modules:
+    # Any module_defaults for testns.testcoll.module will not apply to a module_uses_action_defaults task:
+    #
+    # module_defaults:
+    #   testns.testcoll.module:
+    #     option: value
+    #
+    # But defaults for testns.testcoll.module_uses_action_defaults or testns.testcoll.eos will:
+    #
+    # module_defaults:
+    #   testns.testcoll.module_uses_action_defaults:
+    #     option: value
+    #   testns.testcoll.eos:
+    #     option: defined_last_i_win
+    module_uses_action_defaults:
+      redirect: testns.testcoll.module
+
+    # Not "eos_facts" to ensure TE is not finding handler via prefix
+    # eosfacts tasks should not get eos module_defaults (or defaults for other modules that use eos action plugin)
+    eosfacts:
+      action_plugin: testns.testcoll.eos
+
+    # Test that `action_plugin` has higher precedence than module-redirected-as-action  - reverse this?
+    # Current behavior is iosfacts/ios_facts do not get ios defaults.
+    iosfacts:
+      redirect: testns.testcoll.ios_facts
+    ios_facts:
+      action_plugin: testns.testcoll.redirected_action
+
 action_groups:
   testgroup:
     # Test metadata 'extend_group' feature does not get stuck in a recursive loop

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/action/eos.py

@@ -0,0 +1,18 @@
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import (absolute_import, division, print_function)
+__metaclass__ = type
+
+from ansible.plugins.action.normal import ActionModule as ActionBase
+from ansible.utils.vars import merge_hash
+
+
+class ActionModule(ActionBase):
+
+    def run(self, tmp=None, task_vars=None):
+
+        result = super(ActionModule, self).run(tmp, task_vars)
+        result['action_plugin'] = 'eos'
+
+        return result

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/action/ios.py

@@ -0,0 +1,18 @@
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import (absolute_import, division, print_function)
+__metaclass__ = type
+
+from ansible.plugins.action.normal import ActionModule as ActionBase
+from ansible.utils.vars import merge_hash
+
+
+class ActionModule(ActionBase):
+
+    def run(self, tmp=None, task_vars=None):
+
+        result = super(ActionModule, self).run(tmp, task_vars)
+        result['action_plugin'] = 'ios'
+
+        return result

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/action/vyos.py

@@ -0,0 +1,18 @@
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import (absolute_import, division, print_function)
+__metaclass__ = type
+
+from ansible.plugins.action.normal import ActionModule as ActionBase
+from ansible.utils.vars import merge_hash
+
+
+class ActionModule(ActionBase):
+
+    def run(self, tmp=None, task_vars=None):
+
+        result = super(ActionModule, self).run(tmp, task_vars)
+        result['action_plugin'] = 'vyos'
+
+        return result

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/modules/eosfacts.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+
+DOCUMENTATION = r'''
+---
+module: eosfacts
+short_description: module to test module_defaults
+description: module to test module_defaults
+version_added: '2.13'
+'''
+
+EXAMPLES = r'''
+'''
+
+from ansible.module_utils.basic import AnsibleModule
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            eosfacts=dict(type=bool),
+        ),
+        supports_check_mode=True
+    )
+    module.exit_json(eosfacts=module.params['eosfacts'])
+
+
+if __name__ == '__main__':
+    main()

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/modules/ios_facts.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+
+DOCUMENTATION = r'''
+---
+module: ios_facts
+short_description: module to test module_defaults
+description: module to test module_defaults
+version_added: '2.13'
+'''
+
+EXAMPLES = r'''
+'''
+
+from ansible.module_utils.basic import AnsibleModule
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            ios_facts=dict(type=bool),
+        ),
+        supports_check_mode=True
+    )
+    module.exit_json(ios_facts=module.params['ios_facts'])
+
+
+if __name__ == '__main__':
+    main()

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/modules/module.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+
+DOCUMENTATION = r'''
+---
+module: module
+short_description: module to test module_defaults
+description: module to test module_defaults
+version_added: '2.13'
+'''
+
+EXAMPLES = r'''
+'''
+
+from ansible.module_utils.basic import AnsibleModule
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            action_option=dict(type=bool),
+        ),
+        supports_check_mode=True
+    )
+    module.exit_json(action_option=module.params['action_option'])
+
+
+if __name__ == '__main__':
+    main()

test/integration/targets/module_defaults/collections/ansible_collections/testns/testcoll/plugins/modules/vyosfacts.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+
+# Copyright: (c) 2022, Ansible Project
+# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)
+
+from __future__ import absolute_import, division, print_function
+__metaclass__ = type
+
+
+DOCUMENTATION = r'''
+---
+module: vyosfacts
+short_description: module to test module_defaults
+description: module to test module_defaults
+version_added: '2.13'
+'''
+
+EXAMPLES = r'''
+'''
+
+from ansible.module_utils.basic import AnsibleModule
+
+
+def main():
+    module = AnsibleModule(
+        argument_spec=dict(
+            vyosfacts=dict(type=bool),
+        ),
+        supports_check_mode=True
+    )
+    module.exit_json(vyosfacts=module.params['vyosfacts'])
+
+
+if __name__ == '__main__':
+    main()

test/integration/targets/module_defaults/runme.sh

@@ -2,8 +2,13 @@
 
 set -eux
 
+# Symlink is test for backwards-compat (only workaround for https://github.com/ansible/ansible/issues/77059)
+sudo ln -s "${PWD}/collections/ansible_collections/testns/testcoll/plugins/action/vyos.py" ./collections/ansible_collections/testns/testcoll/plugins/action/vyosfacts.py
+
 ansible-playbook test_defaults.yml "$@"
 
+sudo rm ./collections/ansible_collections/testns/testcoll/plugins/action/vyosfacts.py
+
 ansible-playbook test_action_groups.yml "$@"
 
 ansible-playbook test_action_group_metadata.yml "$@"

test/integration/targets/module_defaults/test_defaults.yml

@@ -110,3 +110,140 @@
         - "builtin_legacy_defaults_2.msg == 'legacy default'"
 
   - include_tasks: tasks/main.yml
+
+- name: test preferring module name defaults for platform-specific actions
+  hosts: localhost
+  gather_facts: no
+  tasks:
+  - name: ensure eosfacts does not use action plugin default
+    testns.testcoll.eosfacts:
+    module_defaults:
+      testns.testcoll.eos:
+        fail: true
+
+  - name: eosfacts does use module name defaults
+    testns.testcoll.eosfacts:
+    module_defaults:
+      testns.testcoll.eosfacts:
+        eosfacts: true
+    register: result
+
+  - assert:
+      that:
+        - result.eosfacts
+        - result.action_plugin == 'eos'
+
+  - name: ensure vyosfacts does not use action plugin default
+    testns.testcoll.vyosfacts:
+    module_defaults:
+      testns.testcoll.vyos:
+        fail: true
+
+  - name: vyosfacts does use vyosfacts defaults
+    testns.testcoll.vyosfacts:
+    module_defaults:
+      testns.testcoll.vyosfacts:
+        vyosfacts: true
+    register: result
+
+  - assert:
+      that:
+        - result.vyosfacts
+        - result.action_plugin == 'vyos'
+
+  - name: iosfacts/ios_facts does not use action plugin default (module action_plugin field has precedence over module-as-action-redirect)
+    collections:
+      - testns.testcoll
+    module_defaults:
+      testns.testcoll.ios:
+        fail: true
+    block:
+      - ios_facts:
+        register: result
+      - assert:
+          that:
+            - result.action_plugin == 'ios'
+
+      - iosfacts:
+        register: result
+      - assert:
+          that:
+            - result.action_plugin == 'ios'
+
+  - name: ensure iosfacts/ios_facts uses ios_facts defaults
+    collections:
+      - testns.testcoll
+    module_defaults:
+      testns.testcoll.ios_facts:
+        ios_facts: true
+    block:
+      - ios_facts:
+        register: result
+      - assert:
+          that:
+            - result.ios_facts
+            - result.action_plugin == 'ios'
+
+      - iosfacts:
+        register: result
+      - assert:
+          that:
+            - result.ios_facts
+            - result.action_plugin == 'ios'
+
+  - name: ensure iosfacts/ios_facts uses iosfacts defaults
+    collections:
+      - testns.testcoll
+    module_defaults:
+      testns.testcoll.iosfacts:
+        ios_facts: true
+    block:
+      - ios_facts:
+        register: result
+      - assert:
+          that:
+            - result.ios_facts
+            - result.action_plugin == 'ios'
+
+      - iosfacts:
+        register: result
+      - assert:
+          that:
+            - result.ios_facts
+            - result.action_plugin == 'ios'
+
+  - name: ensure redirected action gets redirected action defaults
+    testns.testcoll.module_uses_action_defaults:
+    module_defaults:
+      testns.testcoll.module_uses_action_defaults:
+        action_option: true
+    register: result
+
+  - assert:
+      that:
+        - result.action_option
+        - result.action_plugin == 'eos'
+
+  - name: ensure redirected action gets resolved action defaults
+    testns.testcoll.module_uses_action_defaults:
+    module_defaults:
+      testns.testcoll.eos:
+        action_option: true
+    register: result
+
+  - assert:
+      that:
+        - result.action_option
+        - result.action_plugin == 'eos'
+
+  - name: ensure redirected action does not use module-specific defaults
+    testns.testcoll.module_uses_action_defaults:
+    module_defaults:
+      testns.testcoll.module:
+        fail: true
+    register: result
+
+  - assert:
+      that:
+        - not result.action_option
+        - result.action_plugin == 'eos'

test/units/executor/test_task_executor.py

@@ -25,14 +25,18 @@ from units.compat import unittest
 from unittest.mock import patch, MagicMock
 from ansible.errors import AnsibleError
 from ansible.executor.task_executor import TaskExecutor, remove_omit
-from ansible.plugins.loader import action_loader, lookup_loader
+from ansible.plugins.loader import action_loader, lookup_loader, module_loader
 from ansible.parsing.yaml.objects import AnsibleUnicode
 from ansible.utils.unsafe_proxy import AnsibleUnsafeText, AnsibleUnsafeBytes
 from ansible.module_utils.six import text_type
 
+from collections import namedtuple
 from units.mock.loader import DictDataLoader
 
 
+get_with_context_result = namedtuple('get_with_context_result', ['object', 'plugin_load_context'])
+
+
 class TestTaskExecutor(unittest.TestCase):
 
     def test_task_executor_init(self):
@@ -204,6 +208,8 @@ class TestTaskExecutor(unittest.TestCase):
             final_q=MagicMock(),
         )
 
+        context = MagicMock(resolved=False)
+        te._shared_loader_obj.module_loader.find_plugin_with_context.return_value = context
         action_loader = te._shared_loader_obj.action_loader
         action_loader.has_plugin.return_value = True
         action_loader.get.return_value = mock.sentinel.handler
@@ -238,6 +244,8 @@ class TestTaskExecutor(unittest.TestCase):
             final_q=MagicMock(),
         )
 
+        context = MagicMock(resolved=False)
+        te._shared_loader_obj.module_loader.find_plugin_with_context.return_value = context
         action_loader = te._shared_loader_obj.action_loader
         action_loader.has_plugin.side_effect = [False, True]
         action_loader.get.return_value = mock.sentinel.handler
@@ -252,7 +260,7 @@ class TestTaskExecutor(unittest.TestCase):
         handler = te._get_action_handler(mock_connection, mock_templar)
 
         self.assertIs(mock.sentinel.handler, handler)
-        action_loader.has_plugin.assert_has_calls([mock.call(action, collection_list=te._task.collections),
+        action_loader.has_plugin.assert_has_calls([mock.call(action, collection_list=te._task.collections),  # called twice
                                                    mock.call(module_prefix, collection_list=te._task.collections)])
 
         action_loader.get.assert_called_once_with(
@@ -277,6 +285,9 @@ class TestTaskExecutor(unittest.TestCase):
         action_loader.has_plugin.return_value = False
         action_loader.get.return_value = mock.sentinel.handler
         action_loader.__contains__.return_value = False
+        module_loader = te._shared_loader_obj.module_loader
+        context = MagicMock(resolved=False)
+        module_loader.find_plugin_with_context.return_value = context
 
         mock_connection = MagicMock()
         mock_templar = MagicMock()
@@ -302,6 +313,7 @@ class TestTaskExecutor(unittest.TestCase):
         mock_host = MagicMock()
 
         mock_task = MagicMock()
+        mock_task.action = 'mock.action'
         mock_task.args = dict()
         mock_task.retries = 0
         mock_task.delay = -1
@@ -328,7 +340,7 @@ class TestTaskExecutor(unittest.TestCase):
         mock_action = MagicMock()
         mock_queue = MagicMock()
 
-        shared_loader = None
+        shared_loader = MagicMock()
         new_stdin = None
         job_vars = dict(omit="XXXXXXXXXXXXXXXXXXX")
 
@@ -344,7 +356,8 @@ class TestTaskExecutor(unittest.TestCase):
         )
 
         te._get_connection = MagicMock(return_value=mock_connection)
-        te._get_action_handler = MagicMock(return_value=mock_action)
+        context = MagicMock()
+        te._get_action_handler_with_context = MagicMock(return_value=get_with_context_result(mock_action, context))
 
         mock_action.run.return_value = dict(ansible_facts=dict())
         res = te._execute()