tunr copyload_shared_lib into a ctx manager and refactor tests

3 files changed, 48 insertions, 40 deletions

diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index dd95755f..3ccb3044 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -16,6 +16,7 @@ import ctypes
 import errno
 import functools
 import os
+import random
 import re
 import shutil
 import socket
@@ -1019,13 +1020,21 @@ def is_namedtuple(x):
     return all(type(n) == str for n in f)
 
 
-def copyload_shared_lib(src, dst_prefix=TESTFILE_PREFIX):
-    """Given an existing shared so / DLL library copies it in
-    another location and loads it via ctypes.
-    Return the new path.
+@contextlib.contextmanager
+def copyload_shared_lib(dst_prefix=TESTFILE_PREFIX):
+    """Ctx manager which picks up a random shared so/dll lib used
+    by this process, copies it in another location and loads it
+    in memory via ctypes.
+    Return the new absolutized path.
     """
-    newpath = tempfile.mktemp(prefix=dst_prefix,
-                              suffix=os.path.splitext(src)[1])
-    shutil.copyfile(src, newpath)
-    ctypes.CDLL(newpath)
-    return newpath
+    ext = ".so" if POSIX else ".dll"
+    dst = tempfile.mktemp(prefix=dst_prefix, suffix=ext)
+    libs = [x.path for x in psutil.Process().memory_maps()
+            if os.path.normcase(x.path).endswith(ext)]
+    src = random.choice(libs)
+    try:
+        shutil.copyfile(src, dst)
+        ctypes.CDLL(dst)
+        yield os.path.realpath(dst)
+    finally:
+        safe_rmpath(dst)
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 58dc6829..86a784d2 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -656,15 +656,11 @@ class TestProcess(unittest.TestCase):
 
     @unittest.skipIf(not HAS_MEMORY_MAPS, "not supported")
     def test_memory_maps_lists_lib(self):
-        normcase = os.path.normcase
-        realpath = os.path.realpath
-        p = psutil.Process()
-        ext = ".so" if POSIX else ".dll"
-        old = [x.path for x in p.memory_maps()
-               if normcase(x.path).endswith(ext)][0]
-        new = realpath(normcase(copyload_shared_lib(old)))
-        newpaths = [realpath(normcase(x.path)) for x in p.memory_maps()]
-        self.assertIn(new, newpaths)
+        with copyload_shared_lib() as path:
+            # Make sure a newly loaded shared lib is listed.
+            libpaths = [os.path.realpath(os.path.normcase(x.path))
+                        for x in psutil.Process().memory_maps()]
+            self.assertIn(path, libpaths)
 
     def test_memory_percent(self):
         p = psutil.Process()
diff --git a/psutil/tests/test_unicode.py b/psutil/tests/test_unicode.py
index 6dc1e603..2599b2f8 100644
--- a/psutil/tests/test_unicode.py
+++ b/psutil/tests/test_unicode.py
@@ -93,9 +93,9 @@ import psutil
 import psutil.tests
 
 
-def can_deal_with_funky_name(name):
+def subprocess_supports_unicode(name):
     """Return True if both the fs and the subprocess module can
-    deal with a funky file name.
+    deal with a unicode file name.
     """
     if PY3:
         return True
@@ -110,13 +110,22 @@ def can_deal_with_funky_name(name):
         return True
 
 
+def ctypes_supports_unicode(name):
+    if PY3:
+        return True
+    try:
+        with copyload_shared_lib():
+            pass
+    except UnicodeEncodeError:
+        return False
+
+
+# An invalid unicode string.
 if PY3:
     INVALID_NAME = (TESTFN.encode('utf8') + b"f\xc0\x80").decode(
         'utf8', 'surrogateescape')
 else:
     INVALID_NAME = TESTFN + "f\xc0\x80"
-UNICODE_OK = can_deal_with_funky_name(TESTFN_UNICODE)
-INVALID_UNICODE_OK = can_deal_with_funky_name(INVALID_NAME)
 
 
 # ===================================================================
@@ -244,27 +253,21 @@ class _BaseFSAPIsTests(object):
 
     @unittest.skipIf(not HAS_MEMORY_MAPS, "not supported")
     def test_memory_maps(self):
-        normcase = os.path.normcase
-        realpath = os.path.realpath
-        p = psutil.Process()
-        ext = ".so" if POSIX else ".dll"
-        old = [realpath(x.path) for x in p.memory_maps()
-               if normcase(x.path).endswith(ext)][0]
-        try:
-            new = realpath(normcase(
-                copyload_shared_lib(old, dst_prefix=self.funky_name)))
-        except UnicodeEncodeError:
-            if PY3:
-                raise
-            else:
-                raise unittest.SkipTest("ctypes can't handle unicode")
-        newpaths = [realpath(normcase(x.path)) for x in p.memory_maps()]
-        self.assertIn(new, newpaths)
+        if not ctypes_supports_unicode(self.funky_name):
+            raise unittest.SkipTest("ctypes can't handle unicode")
+
+        with copyload_shared_lib(dst_prefix=self.funky_name) as funky_path:
+            libpaths = [os.path.realpath(os.path.normcase(x.path))
+                        for x in psutil.Process().memory_maps()]
+            self.assertIn(funky_path, libpaths)
+            for path in libpaths:
+                self.assertIsInstance(path, str)
 
 
 @unittest.skipIf(OSX and TRAVIS, "unreliable on TRAVIS")  # TODO
 @unittest.skipIf(ASCII_FS, "ASCII fs")
-@unittest.skipIf(not UNICODE_OK, "subprocess can't deal with unicode")
+@unittest.skipIf(not subprocess_supports_unicode(TESTFN_UNICODE),
+                 "subprocess can't deal with unicode")
 class TestFSAPIs(_BaseFSAPIsTests, unittest.TestCase):
     """Test FS APIs with a funky, valid, UTF8 path name."""
     funky_name = TESTFN_UNICODE
@@ -277,7 +280,7 @@ class TestFSAPIs(_BaseFSAPIsTests, unittest.TestCase):
 
 
 @unittest.skipIf(OSX and TRAVIS, "unreliable on TRAVIS")  # TODO
-@unittest.skipIf(not INVALID_UNICODE_OK,
+@unittest.skipIf(not subprocess_supports_unicode(INVALID_NAME),
                  "subprocess can't deal with invalid unicode")
 class TestFSAPIsWithInvalidPath(_BaseFSAPIsTests, unittest.TestCase):
     """Test FS APIs with a funky, invalid path name."""