def test_worker_stdout_stream():
with WorkerWithSocket(("send_and_recv", "stdsocket")) as (send_and_recv, stdsocket):
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
def fun(partition):
ret = [x + 1 for x in partition]
print(ret)
return ret
fun_pickle = toCP(fun)
fun_id = b"fun"
send_and_recv(Commands.SETDATA, fun_id, fun_pickle, expect=Commands.OK)
map_output_id = b"data2"
send_and_recv(Commands.CALL, fun_id, data_id, map_output_id, expect=Commands.OK)
# Poll the worker repeatedly until it's done.
isworking = Commands.FALSE
while isworking != Commands.TRUE:
isworking = send_and_recv(Commands.ISWORKING)
time.sleep(0.1)
received_stdout = stdsocket.recv()
# Ignore the outer worker's logging.
while received_stdout.startswith(b"WORKER"):
received_stdout = stdsocket.recv()
assert received_stdout == str(fun(data)).encode("ascii"), received_stdout
listing = send_and_recv(Commands.LISTDATA)
assert set(fromCP(listing)) == {data_id, fun_id, map_output_id}, listing
send_and_recv(Commands.GETDATA, map_output_id, expect=toCP(fun(data)))
print("") # Make sure this test doesn't end on half a line
def test_worker_del_data():
with WorkerWithSocket() as send_and_recv:
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
send_and_recv(Commands.DELDATA, data_id, expect=Commands.TRUE)
send_and_recv(Commands.DELDATA, data_id, expect=Commands.FALSE)
send_and_recv(Commands.GETDATA, data_id, expect=lambda e: isinstance(fromCP(e), Exception))
def test_worker_set_call_get_data():
with WorkerWithSocket() as send_and_recv:
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
fun_id = b"fun1"
fun = lambda partition: [x + 1 for x in partition]
fun_pickle = toCP(fun)
send_and_recv(Commands.SETDATA, fun_id, fun_pickle, expect=Commands.OK)
map_output_id = b"data2"
send_and_recv(Commands.CALL, fun_id, data_id, map_output_id, expect=Commands.OK)
# Poll the worker repeatedly until it's done.
isworking = Commands.FALSE
while isworking != Commands.TRUE:
isworking = send_and_recv(Commands.ISWORKING)
time.sleep(0.1)
send_and_recv(Commands.LISTDATA, expect=lambda d: set(fromCP(d)) == {data_id, fun_id, map_output_id})
send_and_recv(Commands.GETDATA, map_output_id, expect=toCP(fun(data)))
def get_data(self, data_id):
return [fromCP(r) for r in self.transact_to_all_workers(Commands.GETDATA, itertools.repeat((data_id,)))]
