From 7428d42c16966a042f0b7c8192cd6308a6f4a9f2 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 31 Oct 2021 09:20:06 +0530
Subject: [PATCH] Restore environ after running tests

---
 test.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/test.py b/test.py
index d42415077..d0ed8df9b 100755
--- a/test.py
+++ b/test.py
@@ -7,19 +7,33 @@ import os
 import sys
 import warnings
 from tempfile import TemporaryDirectory
+from contextlib import contextmanager
+from typing import Iterator
 
 base = os.path.dirname(os.path.abspath(__file__))
 
 
+@contextmanager
+def env_vars(**kw: str) -> Iterator[None]:
+    originals = {k: os.environ.get(k) for k in kw}
+    os.environ.update(kw)
+    try:
+        yield
+    finally:
+        for k, v in originals.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v
+
+
 def init_env() -> None:
     sys.path.insert(0, base)
 
 
 def main() -> None:
     warnings.simplefilter('error')
-    os.environ['PYTHONWARNINGS'] = 'error'
-    with TemporaryDirectory() as tdir:
-        os.environ['HOME'] = os.environ['USERPROFILE'] = tdir
+    with TemporaryDirectory() as tdir, env_vars(PYTHONWARNINGS='error', HOME=tdir, USERPROFILE=tdir):
         init_env()
         m = importlib.import_module('kitty_tests.main')
         m.run_tests()  # type: ignore