fix prompt install tests

Author: bckohan

tests/shellcompletion/__init__.py

@@ -93,10 +93,10 @@ def tearDown(self):
         self.remove()
         super().tearDown()
 
-    def verify_install(self, script=None):
+    def verify_install(self, script=None, directory: t.Optional[Path] = None):
         pass
 
-    def verify_remove(self, script=None):
+    def verify_remove(self, script=None, directory: t.Optional[Path] = None):
         pass
 
     def install(
@@ -502,9 +502,11 @@ def test_multi_install(self):
 
     if platform.system() != "Windows":
 
-        def test_prompt_install(self, env={}):
+        def test_prompt_install(self, env={}, directory: t.Optional[Path] = None):
             import pexpect
 
+            env = {**dict(os.environ), **env}
+
             rex = re.compile
             expected = [
                 rex(r"Append the above contents to (?P<file>.*)\?"),  # 0
@@ -521,7 +523,7 @@ def test_prompt_install(self, env={}):
                 "install",
             ]
             self.remove()
-            self.verify_remove()
+            self.verify_remove(directory=directory)
 
             install = pexpect.spawn(self.manage_script, install_command, env=env)
 
@@ -544,7 +546,7 @@ def wait_for_output(child) -> t.Tuple[int, t.Optional[str]]:
                     self.assertEqual(idx, 2)
                     break
 
-            self.verify_remove()
+            self.verify_remove(directory=directory)
 
             # test an install
             install = pexpect.spawn(self.manage_script, install_command, env=env)
@@ -557,4 +559,4 @@ def wait_for_output(child) -> t.Tuple[int, t.Optional[str]]:
                     self.assertEqual(idx, 3)
                     break
 
-            self.verify_install()
+            self.verify_install(directory=directory)

tests/shellcompletion/test_zsh.py

@@ -1,5 +1,6 @@
 import shutil
 from pathlib import Path
+import typing as t
 
 import pytest
 from django.test import TestCase, override_settings
@@ -25,15 +26,21 @@ class ZshTests(_ScriptCompleteTestCase, TestCase):
         f"DJANGO_SETTINGS_MODULE=tests.settings.completion",
     ]
 
-    def verify_install(self, script=None):
+    def verify_install(self, script=None, directory: t.Optional[Path] = None):
         if not script:
             script = self.manage_script
-        self.assertTrue((self.directory / f"_{script}").exists())
+        self.assertTrue((directory / f"_{script}").exists())
 
-    def verify_remove(self, script=None):
+    def verify_remove(self, script=None, directory: t.Optional[Path] = None):
+        directory = directory or self.directory
         if not script:
             script = self.manage_script
-        self.assertFalse((self.directory / f"_{script}").exists())
+        try:
+            self.assertFalse((directory / f"_{script}").exists())
+        except AssertionError:
+            import ipdb
+
+            ipdb.set_trace()
 
 
 @pytest.mark.skipif(shutil.which("zsh") is None, reason="Z-Shell not available")
@@ -65,8 +72,14 @@ def test_no_template_config(self):
 
     if platform.system() != "Windows":
 
-        def test_prompt_install(self, env={}):
+        def test_prompt_install(self, env={}, directory=None):
             zdot_dir = Path(__file__).parent / "zdotdir"
-            zdot_dir.mkdir(exist_ok=True)
-            super().test_prompt_install(env={"ZDOTDIR": str(zdot_dir.absolute())})
-            shutil.rmtree(zdot_dir)
+            try:
+                zdot_dir.mkdir(exist_ok=True)
+                super().test_prompt_install(
+                    env={"ZDOTDIR": str(zdot_dir.absolute())},
+                    directory=zdot_dir / ".zfunc",
+                )
+            finally:
+                if zdot_dir.exists():
+                    shutil.rmtree(zdot_dir)