From 3179951f5c8b3a6e758eefd47f25842e126cd317 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 15 Dec 2018 15:16:32 +0000
Subject: [PATCH] issue #454: fix AttributeError and atexit.yml test.

---
 ansible_mitogen/runner.py                   | 10 ++++-----
 tests/ansible/integration/runner/atexit.yml | 24 ++++++++++-----------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py
index e2b27fba..46539041 100644
--- a/ansible_mitogen/runner.py
+++ b/ansible_mitogen/runner.py
@@ -357,10 +357,10 @@ class Runner(object):
 
 class AtExitWrapper(object):
     """
-    Newer Ansibles use :func:`atexit.register` to trigger tmpdir cleanup when
-    AnsibleModule.tmpdir is responsible for creating its own temporary
-    directory, however with Mitogen processes are preserved across tasks,
-    meaning cleanup must happen earlier.
+    issue #397, #454: Newer Ansibles use :func:`atexit.register` to trigger
+    tmpdir cleanup when AnsibleModule.tmpdir is responsible for creating its
+    own temporary directory, however with Mitogen processes are preserved
+    across tasks, meaning cleanup must happen earlier.
 
     Patch :func:`atexit.register`, catching :func:`shutil.rmtree` calls so they
     can be executed on task completion, rather than on process shutdown.
@@ -401,7 +401,7 @@ class AtExitWrapper(object):
             self.deferred.append((func, targs, kwargs))
             return
 
-        self.original_register(func, *targs, **kwargs)
+        self.original['register'](func, *targs, **kwargs)
 
 
 class ModuleUtilsImporter(object):
diff --git a/tests/ansible/integration/runner/atexit.yml b/tests/ansible/integration/runner/atexit.yml
index 872cdd57..65d27d59 100644
--- a/tests/ansible/integration/runner/atexit.yml
+++ b/tests/ansible/integration/runner/atexit.yml
@@ -1,10 +1,13 @@
-# issue #397: newer Ansibles rely on atexit to cleanup their temporary
-# directories. Ensure atexit handlers run during runner completion.
+# issue #397, #454: newer Ansibles rely on atexit to cleanup their temporary
+# directories. Ensure atexit handlers calling shutil.rmtree() run during runner
+# completion.
 
 - name: integration/runner/atexit.yml
   hosts: test-targets
   gather_facts: false
   any_errors_fatal: false
+  vars:
+    path: /tmp/atexit-should-delete-this
   tasks:
 
     #
@@ -14,18 +17,15 @@
 
     - custom_python_run_script:
         script: |
-          import atexit
-          atexit.register(lambda:
-            open('/tmp/atexit-was-triggered', 'w').write('yep'))
+          import atexit, shutil
+          path = '{{path}}'
+          os.mkdir(path, int('777', 8))
+          atexit.register(shutil.rmtree, path)
 
-    - slurp:
-        path: /tmp/atexit-was-triggered
+    - stat:
+        path: "{{path}}"
       register: out
 
     - assert:
         that:
-        - out.content|b64decode == "yep"
-
-    - file:
-        path: /tmp/atexit-was-triggered
-        state: absent
+        - not out.stat.exists