def test_filtered_communication(s, a, b): c = yield connect(ip=s.ip, port=s.port) f = yield connect(ip=s.ip, port=s.port) yield write(c, {'op': 'register-client', 'client': 'c'}) yield write(f, {'op': 'register-client', 'client': 'f'}) yield read(c) yield read(f) c = BatchedStream(c, 0) f = BatchedStream(f, 0) assert set(s.streams) == {'c', 'f'} yield write(c, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x'))}, 'dependencies': {'x': [], 'y': ['x']}, 'client': 'c', 'keys': ['y']}) yield write(f, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'z': dumps_task((add, 'x', 10))}, 'dependencies': {'x': [], 'z': ['x']}, 'client': 'f', 'keys': ['z']}) msg = yield read(c) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'y' msg = yield read(f) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'z'
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) stream = BatchedStream(stream, 0) stream.send({ 'op': 'update-graph', 'tasks': { 'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x')) }, 'dependencies': { 'x': [], 'y': ['x'] }, 'keys': ['y'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break stream.send({'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed() stream.close()
def test_BatchedStream_raises(): port = 3435 server = MyServer() server.listen(port) client = TCPClient() stream = yield client.connect('127.0.0.1', port) b = BatchedStream(stream, interval=20) stream.close() with pytest.raises(StreamClosedError): yield b.recv() with pytest.raises(StreamClosedError): yield b.send('123')
def test_compute_stream(s, a, b): stream = yield connect(a.ip, a.port) yield write(stream, {'op': 'compute-stream'}) msgs = [{ 'op': 'compute-task', 'function': dumps(inc), 'args': dumps((i, )), 'key': 'x-%d' % i } for i in range(10)] bstream = BatchedStream(stream, 0) for msg in msgs[:5]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' for msg in msgs[5:]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' yield write(stream, {'op': 'close'})
def test_BatchedStream(): port = 3434 server = MyServer() server.listen(port) client = TCPClient() stream = yield client.connect('127.0.0.1', port) b = BatchedStream(stream, interval=20) b.send('hello') b.send('world') result = yield b.recv() assert result == 'hello' result = yield b.recv() assert result == 'hello' result = yield b.recv() assert result == 'world' result = yield b.recv() assert result == 'world' b.close()
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) stream = BatchedStream(stream, 0) stream.send({'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x'))}, 'dependencies': {'x': [], 'y': ['x']}, 'keys': ['y'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break stream.send({'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed() stream.close()
def test_BatchedStream(): port = 3434 server = MyServer() server.listen(port) client = TCPClient() stream = yield client.connect('127.0.0.1', port) b = BatchedStream(stream, interval=20) b.send('hello') b.send('world') result = yield b.recv(); assert result == 'hello' result = yield b.recv(); assert result == 'hello' result = yield b.recv(); assert result == 'world' result = yield b.recv(); assert result == 'world' b.close()
def test_scheduler(s, a, b): stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) stream = BatchedStream(stream, 10) msg = yield read(stream) assert msg['op'] == 'stream-start' # Test update graph yield write(stream, {'op': 'update-graph', 'tasks': valmap(dumps_task, {'x': (inc, 1), 'y': (inc, 'x'), 'z': (inc, 'y')}), 'dependencies': {'x': [], 'y': ['x'], 'z': ['y']}, 'keys': ['x', 'z'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break assert a.data.get('x') == 2 or b.data.get('x') == 2 # Test erring tasks yield write(stream, {'op': 'update-graph', 'tasks': valmap(dumps_task, {'a': (div, 1, 0), 'b': (inc, 'a')}), 'dependencies': {'a': [], 'b': ['a']}, 'keys': ['a', 'b'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'task-erred' and msg['key'] == 'b': break # Test missing data yield write(stream, {'op': 'missing-data', 'keys': ['z']}) s.ensure_idle_ready() while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break # Test missing data without being informed for w in [a, b]: if 'z' in w.data: del w.data['z'] yield write(stream, {'op': 'update-graph', 'tasks': {'zz': dumps_task((inc, 'z'))}, 'dependencies': {'zz': ['z']}, 'keys': ['zz'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'zz': break write(stream, {'op': 'close'}) stream.close()
def test_scheduler(s, a, b): stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) stream = BatchedStream(stream, 10) msg = yield read(stream) assert msg['op'] == 'stream-start' # Test update graph yield write( stream, { 'op': 'update-graph', 'tasks': valmap(dumps_task, { 'x': (inc, 1), 'y': (inc, 'x'), 'z': (inc, 'y') }), 'dependencies': { 'x': [], 'y': ['x'], 'z': ['y'] }, 'keys': ['x', 'z'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break assert a.data.get('x') == 2 or b.data.get('x') == 2 # Test erring tasks yield write( stream, { 'op': 'update-graph', 'tasks': valmap(dumps_task, { 'a': (div, 1, 0), 'b': (inc, 'a') }), 'dependencies': { 'a': [], 'b': ['a'] }, 'keys': ['a', 'b'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'task-erred' and msg['key'] == 'b': break # Test missing data yield write(stream, {'op': 'missing-data', 'keys': ['z']}) s.ensure_idle_ready() while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break # Test missing data without being informed for w in [a, b]: if 'z' in w.data: del w.data['z'] yield write( stream, { 'op': 'update-graph', 'tasks': { 'zz': dumps_task((inc, 'z')) }, 'dependencies': { 'zz': ['z'] }, 'keys': ['zz'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'zz': break write(stream, {'op': 'close'}) stream.close()
def handle_stream(self, stream, address): batched = BatchedStream(stream, interval=10) while True: msg = yield batched.recv() batched.send(msg) batched.send(msg)