import numpy as np
from IPython.parallel import Client
from distarray.client import DistArrayContext
from distarray import odin

import unittest


# To run these tests, you must have an ipcluster running
# For example, `ipcluster -n 4


c = Client()
dv = c[:]
dac = DistArrayContext(dv)


@odin.local(dac)
def local_sin(da):
    return np.sin(da)


@odin.local(dac)
def local_add50(da):
    return da + 50


@odin.local(dac)
def local_sum(da):
    return np.sum(da)


@odin.local(dac)
def local_add_num(da, num):
    return da + num


@odin.local(dac)
def local_add_nums(da, num1, num2, num3):
    return da + num1 + num2 + num3


@odin.local(dac)
def local_add_distarrayproxies(da, dg):
    return da + dg


@odin.local(dac)
def local_add_mixed(da, num1, dg, num2):
    return da + num1 + dg + num2


class TestLocal(unittest.TestCase):

    def setUp(self):
        dv.execute('import numpy as np')
        self.da = dac.empty((1024, 1024))
        self.da.fill(2 * np.pi)

    def test_local_sin(self):
        db = local_sin(self.da)

    def test_local_add(self):
        dc = local_add50(self.da)

    def test_local_sum(self):
        dd = local_sum(self.da)

    def test_local_add_num(self):
        de = local_add_num(self.da, 11)

    def test_local_add_nums(self):
        df = local_add_nums(self.da, 11, 12, 13)

    def test_local_add_distarrayproxies(self):
        dg = dac.empty((1024, 1024))
        dg.fill(33)
        dh = local_add_distarrayproxies(self.da, dg)

    def test_local_add_mixed(self):
        di = dac.empty((1024, 1024))
        di.fill(33)
        dj = local_add_mixed(self.da, 11, di, 12)


if __name__ == '__main__':
    unittest.main()

@@ -795,16 +795,102 @@
  num2%0A%0A%0A
+@odin.local(dac)%0Adef local_add_ndarray(da, num, ndarr):%0A    return da + num + ndarr%0A%0A%0A
 class Te
@@ -1708,16 +1708,202 @@
 i, 12)%0A%0A
+    def test_local_add_ndarray(self):%0A        shp = self.da.get_localshapes()%5B0%5D%0A        ndarr = np.empty(shp)%0A        ndarr.fill(33)%0A        dj = local_add_ndarray(self.da, 11, ndarr)%0A%0A
 %0Aif __na

Add a failing test (trying to add an ndarray locally).