From c148c869e63054af28ef75eaef5a5f308e6efab1 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 1 Nov 2018 20:04:18 +0000
Subject: [PATCH] issue #76, #370: add disconnect cleanup test

---
 ansible_mitogen/services.py                   | 17 +++++++
 .../integration/context_service/all.yml       |  1 +
 .../context_service/disconnect_cleanup.yml    | 46 +++++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 tests/ansible/integration/context_service/disconnect_cleanup.yml

diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py
index 0ce87e09..a5f31b1c 100644
--- a/ansible_mitogen/services.py
+++ b/ansible_mitogen/services.py
@@ -268,6 +268,23 @@ class ContextService(mitogen.service.Service):
         finally:
             self._lock.release()
 
+    @mitogen.service.expose(mitogen.service.AllowParents())
+    def dump(self):
+        """
+        For testing, return a list of dicts describing every currently
+        connected context.
+        """
+        return [
+            {
+                'context_name': context.name,
+                'via': getattr(self._via_by_context.get(context),
+                               'name', None),
+                'refs': self._refs_by_context.get(context),
+            }
+            for context, key in sorted(self._key_by_context.items(),
+                                       key=lambda c_k: c_k[0].context_id)
+        ]
+
     @mitogen.service.expose(mitogen.service.AllowParents())
     def shutdown_all(self):
         """
diff --git a/tests/ansible/integration/context_service/all.yml b/tests/ansible/integration/context_service/all.yml
index e70199f8..c10d67cb 100644
--- a/tests/ansible/integration/context_service/all.yml
+++ b/tests/ansible/integration/context_service/all.yml
@@ -1,2 +1,3 @@
+- import_playbook: disconnect_cleanup.yml
 - import_playbook: lru_one_target.yml
 - import_playbook: reconnection.yml
diff --git a/tests/ansible/integration/context_service/disconnect_cleanup.yml b/tests/ansible/integration/context_service/disconnect_cleanup.yml
new file mode 100644
index 00000000..b657a0dc
--- /dev/null
+++ b/tests/ansible/integration/context_service/disconnect_cleanup.yml
@@ -0,0 +1,46 @@
+# issue #76, #370: ensure context state is forgotten on disconnect, including
+# state of dependent contexts (e.g. sudo, connection delegation, ..).
+
+- name: integration/context_service/disconnect_cleanup.yml
+  hosts: test-targets
+  any_errors_fatal: true
+  tasks:
+  - meta: end_play
+    when: not is_mitogen
+
+  # Start with a clean slate.
+  - mitogen_shutdown_all:
+
+  # Connect a few users.
+  - shell: "true"
+    become: true
+    become_user: "mitogen__user{{item}}"
+    with_items: [1, 2, 3]
+
+  # Verify current state.
+  - mitogen_action_script:
+      script: |
+        self._connection._connect()
+        result['dump'] = self._connection.parent.call_service(
+          service_name='ansible_mitogen.services.ContextService',
+          method_name='dump'
+        )
+    register: out
+
+  - assert:
+      that: out.dump|length == 4  # ssh account + 3 sudo accounts
+
+  - meta: reset_connection
+
+  # Verify current state.
+  - mitogen_action_script:
+      script: |
+        self._connection._connect()
+        result['dump'] = self._connection.parent.call_service(
+          service_name='ansible_mitogen.services.ContextService',
+          method_name='dump'
+        )
+    register: out
+
+  - assert:
+      that: out.dump|length == 1  # just the ssh account