Wrap get_with_context instead of Ansible's own outer get wrapper.

ansible_mitogen/loaders.py

@@ -99,5 +99,5 @@ def assert_supported_release():
 from ansible.plugins.loader import strategy_loader
 
 # These are original, unwrapped implementations
-action_loader__get = action_loader.get
-connection_loader__get = connection_loader.get_with_context
+action_loader__get_with_context = action_loader.get_with_context
+connection_loader__get_with_context = connection_loader.get_with_context

ansible_mitogen/plugins/connection/mitogen_kubectl.py

@@ -48,7 +48,7 @@
 
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'kubectl'
-    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get(
+    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context(
         'kubectl',
         class_only=True,
     )

ansible_mitogen/plugins/connection/mitogen_ssh.py

@@ -60,7 +60,7 @@
 
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'ssh'
-    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get(
+    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context(
         'ssh',
         class_only=True,
     )

ansible_mitogen/strategy.py

@@ -48,6 +48,8 @@
 import ansible.template
 import ansible.utils.sentinel
 
+from ansible.plugins.loader import get_with_context_result
+
 
 def _patch_awx_callback():
     """
@@ -76,12 +78,12 @@ def patch_add_local(self, **kwargs):
 _patch_awx_callback()
 
 
-def wrap_action_loader__get(name, *args, **kwargs):
+def wrap_action_loader__get_with_context(name, *args, **kwargs):
     """
-    While the mitogen strategy is active, trap action_loader.get() calls,
-    augmenting any fetched class with ActionModuleMixin, which replaces various
-    helper methods inherited from ActionBase with implementations that avoid
-    the use of shell fragments wherever possible.
+    While the mitogen strategy is active, trap action_loader.get_with_context()
+    calls, augmenting any fetched class with ActionModuleMixin, which replaces
+    various helper methods inherited from ActionBase with implementations that
+    avoid the use of shell fragments wherever possible.
 
     This is used instead of static subclassing as it generalizes to third party
     action plugins outside the Ansible tree.
@@ -91,13 +93,19 @@ def wrap_action_loader__get(name, *args, **kwargs):
         name = 'mitogen_' + name
     get_kwargs['collection_list'] = kwargs.pop('collection_list', None)
 
-    klass = ansible_mitogen.loaders.action_loader__get(name, **get_kwargs)
+    (klass, context) = ansible_mitogen.loaders.action_loader__get_with_context(
+        name,
+        **get_kwargs
+    )
+
     if klass:
         bases = (ansible_mitogen.mixins.ActionModuleMixin, klass)
         adorned_klass = type(str(name), bases, {})
         if kwargs.get('class_only'):
-            return adorned_klass
-        return adorned_klass(*args, **kwargs)
+            return get_with_context_result(adorned_klass, context)
+        return get_with_context_result(adorned_klass(*args, **kwargs), context)
+
+    return get_with_context_result(None, context)
 
 
 REDIRECTED_CONNECTION_PLUGINS = (
@@ -115,15 +123,16 @@ def wrap_action_loader__get(name, *args, **kwargs):
 )
 
 
-def wrap_connection_loader__get(name, *args, **kwargs):
+def wrap_connection_loader__get_with_context(name, *args, **kwargs):
     """
-    While a Mitogen strategy is active, rewrite connection_loader.get() calls
-    for some transports into requests for a compatible Mitogen transport.
+    While a Mitogen strategy is active, rewrite
+    connection_loader.get_with_context() calls for some transports into
+    requests for a compatible Mitogen transport.
     """
     if name in REDIRECTED_CONNECTION_PLUGINS:
         name = 'mitogen_' + name
 
-    return ansible_mitogen.loaders.connection_loader__get(name, *args, **kwargs)
+    return ansible_mitogen.loaders.connection_loader__get_with_context(name, *args, **kwargs)
 
 
 def wrap_worker__run(self):
@@ -173,8 +182,8 @@ def _install_wrappers(self):
         Install our PluginLoader monkey patches and update global variables
         with references to the real functions.
         """
-        ansible_mitogen.loaders.action_loader.get = wrap_action_loader__get
-        ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get
+        ansible_mitogen.loaders.action_loader.get_with_context = wrap_action_loader__get_with_context
+        ansible_mitogen.loaders.connection_loader.get_with_context = wrap_connection_loader__get_with_context
 
         global worker__run
         worker__run = ansible.executor.process.worker.WorkerProcess.run
@@ -184,11 +193,11 @@ def _remove_wrappers(self):
         """
         Uninstall the PluginLoader monkey patches.
         """
-        ansible_mitogen.loaders.action_loader.get = (
-            ansible_mitogen.loaders.action_loader__get
+        ansible_mitogen.loaders.action_loader.get_with_context = (
+            ansible_mitogen.loaders.action_loader__get_with_context
         )
         ansible_mitogen.loaders.connection_loader.get_with_context = (
-            ansible_mitogen.loaders.connection_loader__get
+            ansible_mitogen.loaders.connection_loader__get_with_context
         )
         ansible.executor.process.worker.WorkerProcess.run = worker__run