import unittest
import distarray as da
from distarray.mpi.mpibase import create_comm_of_size, InvalidCommSizeError


class TestDistributedArrayProtocol(unittest.TestCase):

    def setUp(self):
        try:
            comm = create_comm_of_size(4)
        except InvalidCommSizeError:
            raise unittest.SkipTest('Must run with comm size > 4.')
        else:
            self.larr = da.LocalArray((16,16),
                                     grid_shape=(4,),
                                     comm=comm, buf=None, offset=0)

    def test_has_export(self):
        self.assertTrue(hasattr(self.larr, '__distarray__'))

    def test_export_keys(self):
        required_keys = set(("buffer", "dimdata"))
        export_data = self.larr.__distarray__()
        exported_keys = set(export_data.keys())
        self.assertEqual(required_keys, exported_keys)

    def test_export_buffer(self):
        """See if we actually export a buffer."""
        export_data = self.larr.__distarray__()
        memoryview(export_data['buffer'])

    def test_export_dimdata_len(self):
        """Test if there is a `dimdict` for every dimension."""
        export_data = self.larr.__distarray__()
        dimdata = export_data['dimdata']
        self.assertEqual(len(dimdata), self.larr.ndim)

    def test_export_dimdata_keys(self):
        export_data = self.larr.__distarray__()
        dimdata = export_data['dimdata']
        required_keys = {"disttype", "periodic", "datasize", "gridrank",
                "gridsize", "indices", "blocksize", "padding"}
        for dimdict in dimdata:
            self.assertEqual(required_keys, dimdict.keys())


    @unittest.skip("Import not yet implemented.")
    def test_round_trip(self):
        new_larr = da.fromdap(self.larr)
        self.assertIs(new_larr.local_array, self.larr.local_array)


if __name__ == '__main__':
    try:
        unittest.main()
    except SystemExit:
        pass

@@ -1640,16 +1640,679 @@
 eys())%0A%0A
+    def test_export_dimdata_values(self):%0A        export_data = self.larr.__distarray__()%0A        dimdata = export_data%5B'dimdata'%5D%0A        valid_disttypes = %7BNone, 'b', 'c', 'bc', 'bp', 'u'%7D%0A        for dd in dimdata:%0A            self.assertIn(dd%5B'disttype'%5D, valid_disttypes)%0A            self.assertIsInstance(dd%5B'periodic'%5D, bool)%0A            self.assertIsInstance(dd%5B'datasize'%5D, int)%0A            self.assertIsInstance(dd%5B'gridrank'%5D, int)%0A            self.assertIsInstance(dd%5B'gridsize'%5D, int)%0A            self.assertIsInstance(dd%5B'indices'%5D, slice)%0A            self.assertIsInstance(dd%5B'blocksize'%5D, int)%0A            self.assertEqual(len(dd%5B'padding'%5D), 2)%0A%0A
 %0A    @un

Test exported types in depth.