ansible_mitogen: Fix usage of connection_loader__get.

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

@@ -46,32 +46,25 @@
 import ansible_mitogen.loaders
 
 
-_get_result = ansible_mitogen.loaders.connection_loader__get(
-    'kubectl',
-    class_only=True,
-)
-
-
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'kubectl'
+    (vanilla_class, load_context) = ansible_mitogen.loaders.connection_loader__get_with_context(
+        'kubectl',
+        class_only=True,
+    )
 
     not_supported_msg = (
         'The "mitogen_kubectl" plug-in requires a version of Ansible '
         'that ships with the "kubectl" connection plug-in.'
     )
 
     def __init__(self, *args, **kwargs):
-        if not _get_result:
+        if not Connection.vanilla_class:
             raise ansible.errors.AnsibleConnectionFailure(self.not_supported_msg)
         super(Connection, self).__init__(*args, **kwargs)
 
     def get_extra_args(self):
-        try:
-            # Ansible < 2.10, _get_result is the connection class
-            connection_options = _get_result.connection_options
-        except AttributeError:
-            # Ansible >= 2.10, _get_result is a get_with_context_result
-            connection_options = _get_result.object.connection_options
+        connection_options = Connection.vanilla_class.connection_options
         parameters = []
         for key in connection_options:
             task_var_name = 'ansible_%s' % key

ansible_mitogen/plugins/connection/mitogen_ssh.py

@@ -60,7 +60,7 @@
 
 class Connection(ansible_mitogen.connection.Connection):
     transport = 'ssh'
-    vanilla_class = 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

@@ -47,6 +47,9 @@
 import ansible.executor.process.worker
 import ansible.template
 import ansible.utils.sentinel
+import ansible.playbook.play_context
+
+from ansible.plugins.loader import get_with_context_result
 
 
 def _patch_awx_callback():
@@ -76,12 +79,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 +94,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 +124,26 @@ 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:
+    is_play_using_mitogen_connection = None
+    if len(args) > 0 and isinstance(args[0], ansible.playbook.play_context.PlayContext):
+        play_context = args[0]
+        is_play_using_mitogen_connection = play_context.connection in REDIRECTED_CONNECTION_PLUGINS
+
+    # assume true if we're not in a play context since we're using a Mitogen strategy
+    if is_play_using_mitogen_connection is None:
+        is_play_using_mitogen_connection = True
+
+    redirect_connection = name in REDIRECTED_CONNECTION_PLUGINS and is_play_using_mitogen_connection
+    if redirect_connection:
         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 +193,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 +204,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
 

tests/ansible/regression/all.yml

@@ -13,6 +13,7 @@
 - import_playbook: issue_591__setuptools_cwd_crash.yml
 - import_playbook: issue_615__streaming_transfer.yml
 - import_playbook: issue_655__wait_for_connection_error.yml
+- import_playbook: issue_766__get_with_context.yml
 - import_playbook: issue_776__load_plugins_called_twice.yml
 - import_playbook: issue_952__ask_become_pass.yml
 - import_playbook: issue_1066__add_host__host_key_checking.yml

tests/ansible/regression/issue_766__get_with_context.yml

@@ -0,0 +1,40 @@
+# https://github.com/mitogen-hq/mitogen/issues/776
+---
+- name: regression/issue_766__get_with_context.yml
+  hosts: localhost
+  gather_facts: true
+  tasks:
+    - block:
+        - name: Start container
+          delegate_to: localhost
+          vars:
+            ansible_port: 8040
+          containers.podman.podman_container:
+            name: sysprep
+            image: ghcr.io/mitogen-hq/sysrepo-netopeer2:latest
+            auto_remove: true
+            detach: true
+            recreate: true
+            published_ports:
+              - "{{ ansible_port }}:830"
+
+        - name: Wait for container
+          delegate_to: localhost
+          # TODO robust condition. wait_for + search_regex? wait_for_connection?
+          wait_for:
+            timeout: 5
+
+        - name: Get running configuration and state data
+          vars:
+            ansible_port: 8040
+            ansible_connection: netconf
+            ansible_user: netconf
+            ansible_password: netconf
+          ansible.netcommon.netconf_get:
+
+      always:
+        - name: Cleanup container
+          delegate_to: localhost
+          containers.podman.podman_container:
+            name: sysprep
+            state: absent

tests/ansible/requirements.txt

@@ -3,5 +3,9 @@ paramiko==2.3.2  # Last 2.6-compat version.
 #   ModuleNotFoundError: No module named 'setuptools.command.test'
 # https://github.com/pypa/setuptools/issues/4519
 hdrhistogram==0.6.1
+
+ncclient==0.6.13; python_version <= '2.7'
+ncclient==0.6.16; python_version > '2.7'
+
 PyYAML==3.11; python_version < '2.7'
 PyYAML==5.3.1; python_version >= '2.7'  # Latest release (Jan 2021)