From 20575ca96aaa43440d3583f68266133fa3e65dfc Mon Sep 17 00:00:00 2001
From: "Adam J. Stewart" <ajstewart426@gmail.com>
Date: Tue, 17 Jan 2023 20:47:16 -0700
Subject: Subclass PythonExtension instead of custom test method (#34985)

---
 var/spack/repos/builtin/packages/mxnet/package.py  | 61 +---------------------
 var/spack/repos/builtin/packages/py-dgl/package.py |  9 +---
 .../builtin/packages/py-tensorflow/package.py      | 15 +-----
 3 files changed, 3 insertions(+), 82 deletions(-)

diff --git a/var/spack/repos/builtin/packages/mxnet/package.py b/var/spack/repos/builtin/packages/mxnet/package.py
index ebd0c8f713..a8578d4061 100644
--- a/var/spack/repos/builtin/packages/mxnet/package.py
+++ b/var/spack/repos/builtin/packages/mxnet/package.py
@@ -6,7 +6,7 @@
 from spack.package import *
 
 
-class Mxnet(CMakePackage, CudaPackage):
+class Mxnet(CMakePackage, CudaPackage, PythonExtension):
     """MXNet is a deep learning framework
     designed for both efficiency and flexibility."""
 
@@ -16,51 +16,6 @@ class Mxnet(CMakePackage, CudaPackage):
     git = "https://github.com/apache/incubator-mxnet.git"
 
     maintainers = ["adamjstewart"]
-    import_modules = [
-        "mxnet",
-        "mxnet.numpy_extension",
-        "mxnet.optimizer",
-        "mxnet.module",
-        "mxnet.io",
-        "mxnet.cython",
-        "mxnet.ndarray",
-        "mxnet.gluon",
-        "mxnet.symbol",
-        "mxnet._cy3",
-        "mxnet.contrib",
-        "mxnet.numpy",
-        "mxnet._ffi",
-        "mxnet.image",
-        "mxnet.kvstore",
-        "mxnet.notebook",
-        "mxnet._ctypes",
-        "mxnet.rnn",
-        "mxnet.ndarray.numpy_extension",
-        "mxnet.ndarray.numpy",
-        "mxnet.gluon.nn",
-        "mxnet.gluon.model_zoo",
-        "mxnet.gluon.contrib",
-        "mxnet.gluon.data",
-        "mxnet.gluon.rnn",
-        "mxnet.gluon.model_zoo.vision",
-        "mxnet.gluon.contrib.nn",
-        "mxnet.gluon.contrib.estimator",
-        "mxnet.gluon.contrib.cnn",
-        "mxnet.gluon.contrib.data",
-        "mxnet.gluon.contrib.rnn",
-        "mxnet.gluon.data.vision",
-        "mxnet.symbol.numpy_extension",
-        "mxnet.symbol.numpy",
-        "mxnet.contrib.onnx",
-        "mxnet.contrib.svrg_optimization",
-        "mxnet.contrib.amp",
-        "mxnet.contrib.text",
-        "mxnet.contrib.onnx.mx2onnx",
-        "mxnet.contrib.onnx.onnx2mx",
-        "mxnet.contrib.amp.lists",
-        "mxnet._ffi._cy3",
-        "mxnet._ffi._ctypes",
-    ]
 
     version("master", branch="master", submodules=True)
     version("1.master", branch="v1.x", submodules=True)
@@ -173,17 +128,3 @@ class Mxnet(CMakePackage, CudaPackage):
             with working_dir("python"):
                 args = std_pip_args + ["--prefix=" + prefix, "."]
                 pip(*args)
-
-    def test(self):
-        """Attempts to import modules of the installed package."""
-
-        if "+python" in self.spec:
-            # Make sure we are importing the installed modules,
-            # not the ones in the source directory
-            for module in self.import_modules:
-                self.run_test(
-                    self.spec["python"].command.path,
-                    ["-c", "import {0}".format(module)],
-                    purpose="checking import of {0}".format(module),
-                    work_dir="spack-test",
-                )
diff --git a/var/spack/repos/builtin/packages/py-dgl/package.py b/var/spack/repos/builtin/packages/py-dgl/package.py
index d4e32d2d68..7a17421191 100644
--- a/var/spack/repos/builtin/packages/py-dgl/package.py
+++ b/var/spack/repos/builtin/packages/py-dgl/package.py
@@ -7,7 +7,7 @@
 from spack.package import *
 
 
-class PyDgl(CMakePackage):
+class PyDgl(CMakePackage, PythonExtension):
     """Deep Graph Library (DGL).
 
     DGL is an easy-to-use, high performance and scalable Python package for
@@ -138,10 +138,3 @@ class PyDgl(CMakePackage):
             )
 
         return modules
-
-    @run_after("install")
-    @on_package_attributes(run_tests=True)
-    def import_module_test(self):
-        with working_dir("spack-test", create=True):
-            for module in self.import_modules:
-                python("-c", "import {0}".format(module))
diff --git a/var/spack/repos/builtin/packages/py-tensorflow/package.py b/var/spack/repos/builtin/packages/py-tensorflow/package.py
index b30899797a..af6b073eb7 100644
--- a/var/spack/repos/builtin/packages/py-tensorflow/package.py
+++ b/var/spack/repos/builtin/packages/py-tensorflow/package.py
@@ -10,7 +10,7 @@ from spack.operating_systems.mac_os import macos_version
 from spack.package import *
 
 
-class PyTensorflow(Package, CudaPackage, ROCmPackage):
+class PyTensorflow(Package, CudaPackage, ROCmPackage, PythonExtension):
     """An Open Source Machine Learning Framework for Everyone.
 
     TensorFlow is an end-to-end open source platform for machine learning. It has a
@@ -1037,16 +1037,3 @@ def protobuf_deps():
             args = std_pip_args + ["--prefix=" + prefix, "."]
             pip(*args)
         remove_linked_tree(tmp_path)
-
-    def test(self):
-        """Attempts to import modules of the installed package."""
-
-        # Make sure we are importing the installed modules,
-        # not the ones in the source directory
-        for module in self.import_modules:
-            self.run_test(
-                self.spec["python"].command.path,
-                ["-c", "import {0}".format(module)],
-                purpose="checking import of {0}".format(module),
-                work_dir="spack-test",
-            )
-- 
cgit v1.2.3-70-g09d2