[pypy-commit] pypy py3k: yet another way to set the correct PYTHONPATH, which works also for other test classes

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: py3k
Changeset: r55483:4d17875981b8
Date: 2012-06-07 16:09 +0200
http://bitbucket.org/pypy/pypy/changeset/4d17875981b8/

Log:	yet another way to set the correct PYTHONPATH, which works also for
	other test classes

diff --git a/pypy/translator/goal/test2/test_app_main.py b/pypy/translator/goal/test2/test_app_main.py
--- a/pypy/translator/goal/test2/test_app_main.py
+++ b/pypy/translator/goal/test2/test_app_main.py
@@ -80,6 +80,20 @@
     """)
 
 
+ at contextmanager
+def setpythonpath():
+    old_pythonpath = os.getenv('PYTHONPATH')
+    rootdir = os.path.dirname(autopath.pypydir)
+    os.putenv('PYTHONPATH', rootdir)
+    try:
+        yield
+    finally:
+        if old_pythonpath is None:
+            os.delenv('PYTHONPATH')
+        else:
+            os.putenv('PYTHONPATH', old_pythonpath)
+
+
 class TestParseCommandLine:
 
     def check_options(self, options, sys_argv, expected):
@@ -230,13 +244,8 @@
 
     def spawn(self, argv):
         # make sure that when we do 'import pypy' we get the correct package
-        rootdir = os.path.dirname(autopath.pypydir)
-        old_pythonpath = os.environ['PYTHONPATH']
-        os.environ['PYTHONPATH'] = rootdir
-        try:
+        with setpythonpath():
             return self._spawn(python3, [app_main] + argv)
-        finally:
-            os.environ['PYTHONPATH'] = old_pythonpath
 
     def test_interactive(self):
         child = self.spawn([])
@@ -566,7 +575,8 @@
         return data, process.returncode
 
     def run(self, *args, **kwargs):
-        data, status = self.run_with_status_code(*args, **kwargs)
+        with setpythonpath():
+            data, status = self.run_with_status_code(*args, **kwargs)
         return data
 
     def test_script_on_stdin(self):