1 files changed, 157 insertions, 0 deletions
diff --git a/setuptools/tests/test_msvc9compiler.py b/setuptools/tests/test_msvc9compiler.py
new file mode 100644
index 00000000..970f7679
--- /dev/null
+++ b/setuptools/tests/test_msvc9compiler.py
@@ -0,0 +1,157 @@
+"""msvc9compiler monkey patch test
+
+This test ensures that importing setuptools is sufficient to replace
+the standard find_vcvarsall function with our patched version that
+finds the Visual C++ for Python package.
+"""
+
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+import distutils.errors
+import contextlib
+
+# importing only setuptools should apply the patch
+__import__('setuptools')
+
+class MockReg:
+    """Mock for distutils.msvc9compiler.Reg. We patch it
+    with an instance of this class that mocks out the
+    functions that access the registry.
+    """
+
+    def __init__(self, hkey_local_machine={}, hkey_current_user={}):
+        self.hklm = hkey_local_machine
+        self.hkcu = hkey_current_user
+
+    def __enter__(self):
+        self.original_read_keys = distutils.msvc9compiler.Reg.read_keys
+        self.original_read_values = distutils.msvc9compiler.Reg.read_values
+
+        _winreg = getattr(distutils.msvc9compiler, '_winreg', None)
+        winreg = getattr(distutils.msvc9compiler, 'winreg', _winreg)
+
+        hives = {
+            winreg.HKEY_CURRENT_USER: self.hkcu,
+            winreg.HKEY_LOCAL_MACHINE: self.hklm,
+        }
+
+        def read_keys(cls, base, key):
+            """Return list of registry keys."""
+            hive = hives.get(base, {})
+            return [k.rpartition('\\')[2]
+                    for k in hive if k.startswith(key.lower())]
+
+        def read_values(cls, base, key):
+            """Return dict of registry keys and values."""
+            hive = hives.get(base, {})
+            return dict((k.rpartition('\\')[2], hive[k])
+                        for k in hive if k.startswith(key.lower()))
+
+        distutils.msvc9compiler.Reg.read_keys = classmethod(read_keys)
+        distutils.msvc9compiler.Reg.read_values = classmethod(read_values)
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        distutils.msvc9compiler.Reg.read_keys = self.original_read_keys
+        distutils.msvc9compiler.Reg.read_values = self.original_read_values
+
+@contextlib.contextmanager
+def patch_env(**replacements):
+    """
+    In a context, patch the environment with replacements. Pass None values
+    to clear the values.
+    """
+    saved = dict(
+        (key, os.environ['key'])
+        for key in replacements
+        if key in os.environ
+    )
+
+    # remove values that are null
+    remove = (key for (key, value) in replacements.items() if value is None)
+    for key in list(remove):
+        os.environ.pop(key, None)
+        replacements.pop(key)
+
+    os.environ.update(replacements)
+
+    try:
+        yield saved
+    finally:
+        for key in replacements:
+            os.environ.pop(key, None)
+        os.environ.update(saved)
+
+class TestMSVC9Compiler(unittest.TestCase):
+
+    def test_find_vcvarsall_patch(self):
+        if not hasattr(distutils, 'msvc9compiler'):
+            # skip
+            return
+
+        self.assertEqual(
+            "setuptools.msvc9_support",
+            distutils.msvc9compiler.find_vcvarsall.__module__,
+            "find_vcvarsall was not patched"
+        )
+
+        find_vcvarsall = distutils.msvc9compiler.find_vcvarsall
+        query_vcvarsall = distutils.msvc9compiler.query_vcvarsall
+
+        # No registry entries or environment variable means we should
+        # not find anything
+        with patch_env(VS90COMNTOOLS=None):
+            with MockReg():
+                self.assertIsNone(find_vcvarsall(9.0))
+
+                try:
+                    query_vcvarsall(9.0)
+                    self.fail('Expected DistutilsPlatformError from query_vcvarsall()')
+                except distutils.errors.DistutilsPlatformError:
+                    exc_message = str(sys.exc_info()[1])
+                self.assertIn('aka.ms/vcpython27', exc_message)
+
+        key_32 = r'software\microsoft\devdiv\vcforpython\9.0\installdir'
+        key_64 = r'software\wow6432node\microsoft\devdiv\vcforpython\9.0\installdir'
+
+        # Make two mock files so we can tell whether HCKU entries are
+        # preferred to HKLM entries.
+        mock_installdir_1 = tempfile.mkdtemp()
+        mock_vcvarsall_bat_1 = os.path.join(mock_installdir_1, 'vcvarsall.bat')
+        open(mock_vcvarsall_bat_1, 'w').close()
+        mock_installdir_2 = tempfile.mkdtemp()
+        mock_vcvarsall_bat_2 = os.path.join(mock_installdir_2, 'vcvarsall.bat')
+        open(mock_vcvarsall_bat_2, 'w').close()
+        try:
+            # Ensure we get the current user's setting first
+            with MockReg(
+                hkey_current_user={key_32: mock_installdir_1},
+                hkey_local_machine={
+                    key_32: mock_installdir_2,
+                    key_64: mock_installdir_2,
+                }
+            ):
+                self.assertEqual(mock_vcvarsall_bat_1, find_vcvarsall(9.0))
+
+            # Ensure we get the local machine value if it's there
+            with MockReg(hkey_local_machine={key_32: mock_installdir_2}):
+                self.assertEqual(mock_vcvarsall_bat_2, find_vcvarsall(9.0))
+
+            # Ensure we prefer the 64-bit local machine key
+            # (*not* the Wow6432Node key)
+            with MockReg(
+                hkey_local_machine={
+                    # This *should* only exist on 32-bit machines
+                    key_32: mock_installdir_1,
+                    # This *should* only exist on 64-bit machines
+                    key_64: mock_installdir_2,
+                }
+            ):
+                self.assertEqual(mock_vcvarsall_bat_1, find_vcvarsall(9.0))
+        finally:
+            shutil.rmtree(mock_installdir_1)
+            shutil.rmtree(mock_installdir_2)