From e27237c0bcfe1a3ce0b6de2e34b991997926bd43 Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Sun, 27 Dec 2020 06:24:55 -0600
Subject: Fix import tests for a few packages (#20560)

---
 .../repos/builtin/packages/py-cartopy/package.py   | 24 ++++++++++++++++++++++
 .../repos/builtin/packages/py-ipython/package.py   | 11 ++++++++++
 .../repos/builtin/packages/py-pyzmq/package.py     | 12 +++++++++++
 .../repos/builtin/packages/py-torch/package.py     |  3 +++
 4 files changed, 50 insertions(+)

diff --git a/var/spack/repos/builtin/packages/py-cartopy/package.py b/var/spack/repos/builtin/packages/py-cartopy/package.py
index 8eebd1fb2b..5affefdf9e 100644
--- a/var/spack/repos/builtin/packages/py-cartopy/package.py
+++ b/var/spack/repos/builtin/packages/py-cartopy/package.py
@@ -14,6 +14,12 @@ class PyCartopy(PythonPackage):
 
     maintainers = ['adamjstewart']
 
+    # Skip test files in cartopy.tests
+    import_modules = [
+        'cartopy', 'cartopy.sphinxext', 'cartopy.io', 'cartopy.geodesic',
+        'cartopy.examples', 'cartopy.mpl', 'cartopy.feature'
+    ]
+
     version('0.18.0', sha256='493ced4698361ffabec1a213d2b711dc836117242c304f3b93f5406182fd8bc2')
     version('0.17.0', sha256='137642e63952404ec0841fa0333ad14c58fbbf19cca2a5ac6a38498c4b4998fb')
     version('0.16.0', sha256='cadf62434492c965220b37f0548bc58180466ad6894a1db57dbc51cd43467e5c')
@@ -48,6 +54,24 @@ class PyCartopy(PythonPackage):
 
     phases = ['build_ext', 'install']
 
+    def setup_build_environment(self, env):
+        # Needed for `spack install --test=root py-cartopy`
+        library_dirs = []
+        for dep in self.spec.dependencies(deptype='link'):
+            query = self.spec[dep.name]
+            library_dirs.extend(query.libs.directories)
+
+        # Cartopy uses ctypes.util.find_library, which searches LD_LIBRARY_PATH
+        # Our RPATH logic works fine, but the unit tests fail without this
+        libs = ':'.join(library_dirs)
+        if self.spec.satisfies('platform=darwin'):
+            env.prepend_path('DYLD_FALLBACK_LIBRARY_PATH', libs)
+        else:
+            env.prepend_path('LD_LIBRARY_PATH', libs)
+
+    # Needed for `spack test run py-cartopy`
+    setup_run_environment = setup_build_environment
+
     def build_ext_args(self, spec, prefix):
         args = [
             spec['geos'].headers.include_flags,
diff --git a/var/spack/repos/builtin/packages/py-ipython/package.py b/var/spack/repos/builtin/packages/py-ipython/package.py
index ba62552304..5c7daf55f5 100644
--- a/var/spack/repos/builtin/packages/py-ipython/package.py
+++ b/var/spack/repos/builtin/packages/py-ipython/package.py
@@ -11,6 +11,17 @@ class PyIpython(PythonPackage):
     homepage = "https://pypi.python.org/pypi/ipython"
     url      = "https://pypi.io/packages/source/i/ipython/ipython-7.18.1.tar.gz"
 
+    # IPython.kernel is deprecated and fails to import
+    import_modules = [
+        'IPython', 'IPython.core', 'IPython.core.tests', 'IPython.core.magics',
+        'IPython.sphinxext', 'IPython.terminal',
+        'IPython.terminal.pt_inputhooks', 'IPython.terminal.tests',
+        'IPython.utils', 'IPython.utils.tests', 'IPython.extensions',
+        'IPython.extensions.tests', 'IPython.testing', 'IPython.testing.tests',
+        'IPython.testing.plugin', 'IPython.lib', 'IPython.lib.tests',
+        'IPython.external', 'IPython.external.decorators'
+    ]
+
     version('7.18.1', sha256='a331e78086001931de9424940699691ad49dfb457cea31f5471eae7b78222d5e')
     version('7.5.0',  sha256='e840810029224b56cd0d9e7719dc3b39cf84d577f8ac686547c8ba7a06eeab26')
     version('7.3.0',  sha256='06de667a9e406924f97781bda22d5d76bfb39762b678762d86a466e63f65dc39')
diff --git a/var/spack/repos/builtin/packages/py-pyzmq/package.py b/var/spack/repos/builtin/packages/py-pyzmq/package.py
index 3e5dae326a..901ca58aec 100644
--- a/var/spack/repos/builtin/packages/py-pyzmq/package.py
+++ b/var/spack/repos/builtin/packages/py-pyzmq/package.py
@@ -32,3 +32,15 @@ class PyPyzmq(PythonPackage):
     depends_on('py-cffi', type=('build', 'run'))
     depends_on('py-gevent', type=('build', 'run'))
     depends_on('libzmq')
+
+    def setup_build_environment(self, env):
+        # Needed for `spack install --test=root py-pyzmq`
+        # Fixes import failure for zmq.backend.cffi
+        # https://github.com/zeromq/pyzmq/issues/395#issuecomment-22041019
+        env.prepend_path(
+            'C_INCLUDE_PATH', self.spec['libzmq'].headers.directories[0])
+        env.prepend_path(
+            'LIBRARY_PATH', self.spec['libzmq'].libs.directories[0])
+
+    # Needed for `spack test run py-pyzmq`
+    setup_run_environment = setup_build_environment
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 7022e562ff..19991da898 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -16,6 +16,9 @@ class PyTorch(PythonPackage, CudaPackage):
 
     maintainers = ['adamjstewart']
 
+    # Exact set of modules is version- and variant-specific
+    import_modules = ['torch', 'torch.autograd', 'torch.nn', 'torch.utils']
+
     version('master', branch='master', submodules=True)
     version('1.7.1', tag='v1.7.1', submodules=True)
     version('1.7.0', tag='v1.7.0', submodules=True)
-- 
cgit v1.2.3-60-g2f50