Merge pull request #402 from enthought/bugfix/389_bobs-version

Add a ContextTestCase to testing.py

Author: kwmsmith

distarray/dist/distarray.py

@@ -42,7 +42,7 @@ def __init__(self, distribution, dtype=float):
         # FIXME: code duplication with context.py.
         ctx = distribution.context
         # FIXME: this is bad...
-        comm_name = ctx.comm
+        comm_name = distribution.comm
         # FIXME: and this is bad...
         da_key = ctx._generate_key()
         ddpr = distribution.get_dim_data_per_rank()

distarray/dist/maps.py

@@ -580,7 +580,7 @@ def __init__(self, context, global_dim_data, targets=None):
         self.grid_shape = tuple(m.grid_size for m in self.maps)
 
         self.grid_shape = normalize_grid_shape(self.grid_shape, self.ndim,
-                                               self.dist, len(context.targets))
+                                               self.dist, len(self.targets))
 
         nelts = reduce(operator.mul, self.grid_shape, 1)
         self.rank_from_coords = np.arange(nelts).reshape(self.grid_shape)

distarray/dist/tests/test_context.py

@@ -17,26 +17,22 @@
 
 import numpy
 
+from distarray.testing import ContextTestCase, check_targets
 from distarray.dist.context import Context
 from distarray.dist.maps import Distribution
 from distarray.dist.ipython_utils import IPythonClient
 from distarray.local import LocalArray
 
 
-class TestContext(unittest.TestCase):
+class TestContext(ContextTestCase):
     """Test Context methods"""
 
     @classmethod
     def setUpClass(cls):
-        cls.context = Context()
+        super(TestContext, cls).setUpClass()
         cls.ndarr = numpy.arange(16).reshape(4, 4)
         cls.darr = cls.context.fromndarray(cls.ndarr)
 
-    @classmethod
-    def tearDownClass(cls):
-        """Close the client connections"""
-        cls.context.close()
-
     def test_get_localarrays(self):
         las = self.darr.get_localarrays()
         self.assertIsInstance(las[0], LocalArray)
@@ -49,68 +45,61 @@ def test_get_ndarrays(self):
 class TestContextCreation(unittest.TestCase):
     """Test Context Creation"""
 
+    @classmethod
+    def setUpClass(cls):
+        cls.client = IPythonClient()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.client.close()
+
     def test_create_Context(self):
         """Can we create a plain vanilla context?"""
-        client = IPythonClient()
-        dac = Context(client)
-        self.assertIs(dac.client, client)
-        dac.close()
-        client.close()
+        dac = Context(self.client)
+        self.assertIs(dac.client, self.client)
 
     def test_create_Context_with_targets(self):
         """Can we create a context with a subset of engines?"""
-        client = IPythonClient()
-        dac = Context(client, targets=[0, 1])
-        self.assertIs(dac.client, client)
+        check_targets(required=2, available=len(self.client))
+        dac = Context(self.client, targets=[0, 1])
+        self.assertIs(dac.client, self.client)
         dac.close()
-        client.close()
 
     def test_create_Context_with_targets_ranks(self):
         """Check that the target <=> rank mapping is consistent."""
-        client = IPythonClient()
+        check_targets(required=4, available=len(self.client))
         targets = [3, 2]
-        dac = Context(client, targets=targets)
+        dac = Context(self.client, targets=targets)
         self.assertEqual(set(dac.targets), set(targets))
         dac.close()
-        client.close()
 
     def test_context_target_reordering(self):
         """Are contexts' targets reordered in a consistent way?"""
-        client = IPythonClient()
-        orig_targets = client.ids
+        orig_targets = self.client.ids
         targets1 = orig_targets[:]
         targets2 = orig_targets[:]
         shuffle(targets1)
         shuffle(targets2)
-        ctx1 = Context(client, targets=targets1)
-        ctx2 = Context(client, targets=targets2)
+        ctx1 = Context(self.client, targets=targets1)
+        ctx2 = Context(self.client, targets=targets2)
         self.assertEqual(ctx1.targets, ctx2.targets)
         ctx1.close()
         ctx2.close()
-        client.close()
 
     def test_create_delete_key(self):
         """ Check that a key can be created and then destroyed. """
-        client = IPythonClient()
-        dac = Context(client)
+        dac = Context(self.client)
         # Create and push a key/value.
         key, value = dac._generate_key(), 'test'
         dac._push({key: value}, targets=dac.targets)
         # Delete the key.
         dac.delete_key(key)
         dac.close()
-        client.close()
-
 
-class TestPrimeCluster(unittest.TestCase):
 
-    @classmethod
-    def setUpClass(cls):
-        cls.context = Context(targets=range(3))
+class TestPrimeCluster(ContextTestCase):
 
-    @classmethod
-    def tearDownClass(cls):
-        cls.context.close()
+    ntargets = 3
 
     def test_1D(self):
         d = Distribution.from_shape(self.context, (3,))
@@ -139,12 +128,9 @@ def test_3D(self):
         self.assertEqual(c.grid_shape, (1, 1, 3))
 
 
-class TestApply(unittest.TestCase):
+class TestApply(ContextTestCase):
 
-    @classmethod
-    def setUpClass(cls):
-        cls.context = Context()
-        cls.num_targets = len(cls.context.targets)
+    ntargets = 'any'
 
     def test_apply_no_args(self):
 
@@ -153,7 +139,7 @@ def foo():
 
         val = self.context.apply(foo)
 
-        self.assertEqual(val, [42] * self.num_targets)
+        self.assertEqual(val, [42] * self.ntargets)
 
     def test_apply_pos_args(self):
 
@@ -162,19 +148,19 @@ def foo(a, b, c):
 
         # push all arguments
         val = self.context.apply(foo, (1, 2, 3))
-        self.assertEqual(val, [6] * self.num_targets)
+        self.assertEqual(val, [6] * self.ntargets)
 
         # some local, some pushed
         local_thing = self.context._key_and_push(2)[0]
         val = self.context.apply(foo, (1, local_thing, 3))
 
-        self.assertEqual(val, [6] * self.num_targets)
+        self.assertEqual(val, [6] * self.ntargets)
 
         # all pushed
         local_args = self.context._key_and_push(1, 2, 3)
         val = self.context.apply(foo, local_args)
 
-        self.assertEqual(val, [6] * self.num_targets)
+        self.assertEqual(val, [6] * self.ntargets)
 
     def test_apply_kwargs(self):
 
@@ -186,25 +172,25 @@ def foo(a, b, c=None, d=None):
         # empty kwargs
         val = self.context.apply(foo, (1, 2))
 
-        self.assertEqual(val, [0] * self.num_targets)
+        self.assertEqual(val, [0] * self.ntargets)
 
         # some empty
         val = self.context.apply(foo, (1, 2), {'d': 3})
 
-        self.assertEqual(val, [5] * self.num_targets)
+        self.assertEqual(val, [5] * self.ntargets)
 
         # all kwargs
         val = self.context.apply(foo, (1, 2), {'c': 2, 'd': 3})
 
-        self.assertEqual(val, [8] * self.num_targets)
+        self.assertEqual(val, [8] * self.ntargets)
 
         # now with local values
         local_a = self.context._key_and_push(1)[0]
         local_c = self.context._key_and_push(3)[0]
 
         val = self.context.apply(foo, (local_a, 2), {'c': local_c, 'd': 3})
 
-        self.assertEqual(val, [9] * self.num_targets)
+        self.assertEqual(val, [9] * self.ntargets)
 
     def test_apply_proxyize(self):
 
@@ -229,7 +215,7 @@ def bar(obj):
             return obj + 10
         val = self.context.apply(bar, (name,))
 
-        self.assertEqual(val, [20] * self.num_targets)
+        self.assertEqual(val, [20] * self.ntargets)
 
     def test_apply_proxyize_sync(self):
 

distarray/dist/tests/test_decorators.py

@@ -17,6 +17,7 @@
 import numpy
 from numpy.testing import assert_array_equal
 
+from distarray.testing import ContextTestCase, check_targets
 from distarray.dist.context import Context
 from distarray.dist.maps import Distribution
 from distarray.dist.decorators import DecoratorBase, local, vectorize
@@ -78,7 +79,9 @@ def dummy_func(*args, **kwargs):
         self.assertTrue(db.key in kw_keys2)
 
 
-class TestLocalDecorator(TestCase):
+class TestLocalDecorator(ContextTestCase):
+
+    ntargets = 'any'
 
     # Functions for @local decorator tests. These are here so we can
     # guarantee they are pushed to the engines before we try to use them.
@@ -143,15 +146,11 @@ def parameterless():
 
     @classmethod
     def setUpClass(cls):
-        cls.context = Context()
+        super(TestLocalDecorator, cls).setUpClass()
         distribution = Distribution.from_shape(cls.context, (5, 5))
         cls.da = cls.context.empty(distribution)
         cls.da.fill(2 * numpy.pi)
 
-    @classmethod
-    def tearDownClass(cls):
-        cls.context.close()
-
     def test_local(self):
         """Test the @local decorator"""
         context = Context()
@@ -178,6 +177,8 @@ def fill_da(da):
         assert_array_equal(da.toarray(), a)
 
     def test_different_contexts(self):
+        check_targets(required=4, available=len(self.context.targets))
+
         ctx1 = Context(targets=range(4))
         ctx2 = Context(targets=range(3))
         distribution1 = Distribution.from_shape(ctx1, (10,))
@@ -200,6 +201,8 @@ def test_local_add50(self):
 
     def test_local_sum(self):
         dd = self.local_sum(self.da)
+        if self.ntargets == 1:
+            dd = [dd]
         lshapes = self.da.get_localshapes()
         expected = []
         for lshape in lshapes: