From 7080751f13ca2bbde8efb200046052b30716bc1a Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Wed, 28 Feb 2018 01:42:37 +0545
Subject: [PATCH] ansible: support environment: too.

---
 ansible_mitogen/helpers.py        | 18 ++++++++++++++----
 ansible_mitogen/mixins.py         |  4 ++++
 examples/playbook/environment.yml | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 4 deletions(-)
 create mode 100644 examples/playbook/environment.yml

diff --git a/ansible_mitogen/helpers.py b/ansible_mitogen/helpers.py
index df2b32f7..2bd0b28c 100644
--- a/ansible_mitogen/helpers.py
+++ b/ansible_mitogen/helpers.py
@@ -96,7 +96,7 @@ def monkey_fail_json(self, **kwargs):
     raise ModuleError(kwargs.get('msg'), kwargs)
 
 
-def run_module(module, raw_params=None, args=None):
+def run_module(module, raw_params=None, args=None, env=None):
     """
     Set up the process environment in preparation for running an Ansible
     module. This monkey-patches the Ansible libraries in various places to
@@ -114,6 +114,10 @@ def run_module(module, raw_params=None, args=None):
         'ANSIBLE_MODULE_ARGS': args
     })
 
+    if env:
+        original_env = os.environ.copy()
+        os.environ.update((k, str(v)) for k, v in env.iteritems())
+
     try:
         mod = __import__(module, {}, {}, [''])
         # Ansible modules begin execution on import. Thus the above __import__
@@ -123,16 +127,22 @@ def run_module(module, raw_params=None, args=None):
         # explicitly.
         mod.main()
     except (Exit, ModuleError), e:
-        return json.dumps(e.dct)
+        result = json.dumps(e.dct)
+
+    if env:
+        os.environ.clear()
+        os.environ.update(original_env)
+
+    return result
 
 
-def _async_main(job_id, module, raw_params, args):
+def _async_main(job_id, module, raw_params, args, env):
     """
     Implementation for the thread that implements asynchronous module
     execution.
     """
     try:
-        rc = run_module(module, raw_params, args)
+        rc = run_module(module, raw_params, args, env)
     except Exception, e:
         rc = mitogen.core.CallError(e)
 
diff --git a/ansible_mitogen/mixins.py b/ansible_mitogen/mixins.py
index 9ac66a36..ddcc0261 100644
--- a/ansible_mitogen/mixins.py
+++ b/ansible_mitogen/mixins.py
@@ -165,11 +165,15 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase):
         else:
             helper = ansible_mitogen.helpers.run_module
 
+        env = {}
+        self._compute_environment_string(env)
+
         # replaces 110 lines
         js = self.call(
             helper,
             get_command_module_name(module_name),
             args=cast(module_args),
+            env=cast(env),
         )
 
         data = self._parse_returned_data({
diff --git a/examples/playbook/environment.yml b/examples/playbook/environment.yml
new file mode 100644
index 00000000..c24bf083
--- /dev/null
+++ b/examples/playbook/environment.yml
@@ -0,0 +1,17 @@
+---
+# Ensure environment: is preserved during call.
+
+- hosts: all
+  gather_facts: false
+  tasks:
+
+    - shell: echo $SOME_ENV
+      environment:
+        SOME_ENV: 123
+      register: result
+
+    - debug: msg={{result}}
+
+    - assert:
+        that: "result.stdout == '123'"
+