# Path: taskloaf/mpi.py
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()
# 
# Path: taskloaf/test_decorators.py
# def mpi_procs(n):
#     n_procs_available = MPI.COMM_WORLD.Get_size()
# 
#     def decorator(test_fnc):
#         test_fnc.n_procs = n
#         return test_fnc
# 
#     return decorator


import cloudpickle
from multiprocessing import Process, Queue
from mpi4py import MPI
from taskloaf.local import LocalComm
from taskloaf.mpi import MPIComm
from taskloaf.zmq import ZMQComm, zmqrun
from taskloaf.test_decorators import mpi_procs

test_cfg = dict()

def queue_test_helper(q):
    q.put(123)

def test_queue():
    q = Queue()
    p = Process(target = queue_test_helper, args = (q, ))
    p.start()
    assert(q.get() == 123)
    p.join()

def test_local_comm():
    qs = [Queue(), Queue()]
    c0 = LocalComm(qs, 0)
    c1 = LocalComm(qs, 1)
    c0.send(1, cloudpickle.dumps(456))
    go = True
    while go:
        val = c1.recv()
        if val is not None:
            assert(cloudpickle.loads(val) == 456)
            go = False

def test_empty_recv():

# Path: taskloaf/mpi.py
# class MPIComm:
#     default_comm = MPI.COMM_WORLD
# 
#     # TODO: multi part messages can be done with tags.
#     def __init__(self, comm=None):
#         if comm is None:
#             comm = MPIComm.default_comm
#         self.comm = comm
#         self.addr = rank(self.comm)
# 
#     def send(self, to_addr, data):
#         self.comm.Isend(data, dest=to_addr)
# 
#     def recv(self):
#         s = MPI.Status()
#         msg_exists = self.comm.iprobe(status=s)
#         if not msg_exists:
#             return None
#         out = memoryview(bytearray(s.count))
#         self.comm.Recv(out, source=s.source)
#         return out
# 
#     def barrier(self):
#         self.comm.Barrier()
# 
# Path: taskloaf/test_decorators.py
# def mpi_procs(n):
#     n_procs_available = MPI.COMM_WORLD.Get_size()
# 
#     def decorator(test_fnc):
#         test_fnc.n_procs = n
#         return test_fnc
# 
#     return decorator


# Path: tests/test_comms.py
import cloudpickle
from multiprocessing import Process, Queue
from mpi4py import MPI
from taskloaf.local import LocalComm
from taskloaf.mpi import MPIComm
from taskloaf.zmq import ZMQComm, zmqrun
from taskloaf.test_decorators import mpi_procs


test_cfg = dict()

def queue_test_helper(q):
    q.put(123)

def test_queue():
    q = Queue()
    p = Process(target = queue_test_helper, args = (q, ))
    p.start()
    assert(q.get() == 123)
    p.join()

def test_local_comm():
    qs = [Queue(), Queue()]
    c0 = LocalComm(qs, 0)
    c1 = LocalComm(qs, 1)
    c0.send(1, cloudpickle.dumps(456))
    go = True
    while go:
        val = c1.recv()
        if val is not None:
            assert(cloudpickle.loads(val) == 456)
            go = False

def test_empty_recv():

    c = MPIComm()