Set backends in setUp and tearDown of CubTests

asi1024 · asi1024 · commit 43b2da6fb720 · 2020-06-29T06:45:19.000Z
diff --git a/cupy/cuda/__init__.py b/cupy/cuda/__init__.py
@@ -35,6 +35,12 @@
 except ImportError:
     thrust_enabled = False
 
+try:
+    from cupy.cuda import cub  # NOQA
+    cub_enabled = True
+except ImportError:
+    cub_enabled = False
+
 try:
     from cupy.cuda import nccl  # NOQA
     nccl_enabled = True
diff --git a/tests/cupy_tests/core_tests/test_ndarray_reduction.py b/tests/cupy_tests/core_tests/test_ndarray_reduction.py
@@ -216,10 +216,16 @@ def test_ptp_nan_imag(self, xp, dtype):
     'order': ('C', 'F'),
 }))
 @testing.gpu
-@unittest.skipUnless(
-    cupy.core._backend.get_routine_backends() == ['cub'],
-    'The CUB routine is not enabled')
+@unittest.skipUnless(cupy.cuda.cub_enabled, 'The CUB routine is not enabled')
 class TestCubReduction(unittest.TestCase):
+
+    def setUp(self):
+        self.old_backends = cupy.core._backend.get_routine_backends()
+        cupy.core._backend.set_routine_backends(['cub'])
+
+    def tearDown(self):
+        cupy.core._backend.set_routine_backends(self.old_backends)
+
     @testing.for_contiguous_axes()
     @testing.for_all_dtypes(no_bool=True, no_float16=True)
     @testing.numpy_cupy_allclose(rtol=1E-5)
diff --git a/tests/cupy_tests/math_tests/test_sumprod.py b/tests/cupy_tests/math_tests/test_sumprod.py
@@ -198,10 +198,16 @@ def test_prod_dtype(self, xp, src_dtype, dst_dtype):
     'order': ('C', 'F'),
 }))
 @testing.gpu
-@unittest.skipUnless(
-    cupy.core._backend.get_routine_backends() == ['cub'],
-    'The CUB routine is not enabled')
-class TestCUBreduction(unittest.TestCase):
+@unittest.skipUnless(cupy.cuda.cub_enabled, 'The CUB routine is not enabled')
+class TestCubReduction(unittest.TestCase):
+
+    def setUp(self):
+        self.old_backends = cupy.core._backend.get_routine_backends()
+        cupy.core._backend.set_routine_backends(['cub'])
+
+    def tearDown(self):
+        cupy.core._backend.set_routine_backends(self.old_backends)
+
     @testing.for_contiguous_axes()
     # sum supports less dtypes; don't test float16 as it's not as accurate?
     @testing.for_dtypes('lLfdFD')
diff --git a/tests/cupy_tests/sorting_tests/test_search.py b/tests/cupy_tests/sorting_tests/test_search.py
@@ -166,10 +166,16 @@ def test_argmin_zero_size_axis1(self, xp, dtype):
     'order': ('C', 'F'),
 }))
 @testing.gpu
-@unittest.skipUnless(
-    cupy.core._backend.get_routine_backends() == ['cub'],
-    'The CUB routine is not enabled')
-class TestCUBreduction(unittest.TestCase):
+@unittest.skipUnless(cupy.cuda.cub_enabled, 'The CUB routine is not enabled')
+class TestCubReduction(unittest.TestCase):
+
+    def setUp(self):
+        self.old_backends = cupy.core._backend.get_routine_backends()
+        cupy.core._backend.set_routine_backends(['cub'])
+
+    def tearDown(self):
+        cupy.core._backend.set_routine_backends(self.old_backends)
+
     @testing.for_dtypes('bhilBHILefdFD')
     @testing.numpy_cupy_allclose(rtol=1E-5)
     def test_cub_argmin(self, xp, dtype):
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
@@ -1534,10 +1534,16 @@ def test_getitem_slice_stop_too_large(self, xp, sp):
 }))
 @testing.with_requires('scipy')
 @testing.gpu
-@unittest.skipUnless(
-    cupy.core._backend.get_routine_backends() == ['cub'],
-    'The CUB routine is not enabled')
-class TestCUBspmv(unittest.TestCase):
+@unittest.skipUnless(cupy.cuda.cub_enabled, 'The CUB routine is not enabled')
+class TestCubSpmv(unittest.TestCase):
+
+    def setUp(self):
+        self.old_backends = cupy.core._backend.get_routine_backends()
+        cupy.core._backend.set_routine_backends(['cub'])
+
+    def tearDown(self):
+        cupy.core._backend.set_routine_backends(self.old_backends)
+
     @property
     def make(self):
         return globals()[self.make_method]
