From 15d68b3c326d0150fbdcc5b811f6fe64e4ad4a52 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 19 Jul 2018 10:58:38 -0400
Subject: [PATCH] issue #309: fix environment cleanup regression.

Closes #309.
---
 ansible_mitogen/runner.py                     | 10 ++--
 docs/changelog.rst                            |  5 ++
 tests/ansible/integration/runner/all.yml      |  1 +
 .../runner/environment_isolation.yml          | 50 +++++++++++++++++++
 .../modules/custom_python_modify_environ.py   | 22 ++++++++
 5 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 tests/ansible/integration/runner/environment_isolation.yml
 create mode 100644 tests/ansible/lib/modules/custom_python_modify_environ.py

diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py
index a8384e21..3960cd38 100644
--- a/ansible_mitogen/runner.py
+++ b/ansible_mitogen/runner.py
@@ -288,9 +288,13 @@ class TemporaryEnvironment(object):
                 os.environ[key] = str(value)
 
     def revert(self):
-        if self.env:
-            os.environ.clear()
-            os.environ.update(self.original)
+        """
+        Revert changes made by the module to the process environment. This must
+        always run, as some modules (e.g. git.py) set variables like GIT_SSH
+        that must be cleared out between runs.
+        """
+        os.environ.clear()
+        os.environ.update(self.original)
 
 
 class TemporaryArgv(object):
diff --git a/docs/changelog.rst b/docs/changelog.rst
index e853f6e8..d0d054c6 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -24,6 +24,11 @@ Mitogen for Ansible
 * `#299 <https://github.com/dw/mitogen/pull/299>`_: fix the ``network_cli``
   connection type when the Mitogen strategy is active.
 
+* `#309 <https://github.com/dw/mitogen/pull/309>`_: fix a regression to process
+  environment cleanup, caused by the change in v0.2.1 to run local tasks with
+  the correct environment.
+
+
 Core Library
 ~~~~~~~~~~~~
 
diff --git a/tests/ansible/integration/runner/all.yml b/tests/ansible/integration/runner/all.yml
index a5da1c87..eaeb46d5 100644
--- a/tests/ansible/integration/runner/all.yml
+++ b/tests/ansible/integration/runner/all.yml
@@ -11,4 +11,5 @@
 - import_playbook: custom_python_new_style_module.yml
 - import_playbook: custom_python_want_json_module.yml
 - import_playbook: custom_script_interpreter.yml
+- import_playbook: environment_isolation.yml
 - import_playbook: forking_behaviour.yml
diff --git a/tests/ansible/integration/runner/environment_isolation.yml b/tests/ansible/integration/runner/environment_isolation.yml
new file mode 100644
index 00000000..08f0924f
--- /dev/null
+++ b/tests/ansible/integration/runner/environment_isolation.yml
@@ -0,0 +1,50 @@
+# issue #309: ensure process environment is restored after a module runs.
+
+
+- name: integration/runner/environment_isolation.yml
+  hosts: test-targets
+  any_errors_fatal: true
+  gather_facts: true
+  tasks:
+
+    # ---
+    # Verify custom env setting is cleared out.
+    # ---
+
+    # Verify sane state first.
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: not out.env.evil_key is defined
+
+    - shell: echo 'hi'
+      environment:
+        evil_key: evil
+
+    # Verify environment was cleaned up.
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: not out.env.evil_key is defined
+
+
+    # ---
+    # Verify non-explicit module env mutations are cleared out.
+    # ---
+
+    # Verify sane state first.
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: not out.env.evil_key is defined
+
+    - custom_python_modify_environ:
+        key: evil_key
+        val: evil
+
+    # Verify environment was cleaned up.
+    - custom_python_detect_environment:
+      register: out
+    - assert:
+        that: not out.env.evil_key is defined
+
diff --git a/tests/ansible/lib/modules/custom_python_modify_environ.py b/tests/ansible/lib/modules/custom_python_modify_environ.py
new file mode 100644
index 00000000..8cdd3bde
--- /dev/null
+++ b/tests/ansible/lib/modules/custom_python_modify_environ.py
@@ -0,0 +1,22 @@
+#!/usr/bin/python
+# I am an Ansible new-style Python module. I modify the process environment and
+# don't clean up after myself.
+
+from ansible.module_utils.basic import AnsibleModule
+
+import os
+import pwd
+import socket
+import sys
+
+
+def main():
+    module = AnsibleModule(argument_spec={
+        'key': {'type': str},
+        'val': {'type': str}
+    })
+    os.environ[module.params['key']] = module.params['val']
+    module.exit_json(msg='Muahahaha!')
+
+if __name__ == '__main__':
+    main()