async def test_push_pull_max_outbox_size(event_loop, endpoint): async with azmq.Context() as context: push_socket = context.socket(azmq.PUSH) push_socket.max_outbox_size = 1 pull_socket = context.socket(azmq.PULL) try: # The PUSH socket connects before the PULL sockets binds and sends # two messages right away. The first one should not block but the # second should as the queue size is 1 and the socket can't # possibly have sent the element so far. push_socket.connect(endpoint) await fivesec(push_socket.send_multipart([b'one'])) await zerosec(push_socket.send_multipart([b'two'])) # The PULL sockets finally binds, and should receive the first # message, even late. pull_socket.bind(endpoint) message = await fivesec(pull_socket.recv_multipart()) assert message == [b'one'] await zerosec(pull_socket.recv_multipart()) finally: push_socket.close() pull_socket.close()
async def test_incompatible_mechanisms(event_loop): async with azmq.Context() as context: c_public_key, c_secret_key = curve_gen_keypair() s_public_key, s_secret_key = curve_gen_keypair() authenticator = ZAPAuthenticator(context) authenticator.add_authorized_key(key=c_public_key) context.set_zap_authenticator(authenticator) req_socket = context.socket( azmq.REQ, mechanism=CurveClient( public_key=c_public_key, secret_key=c_secret_key, server_key=s_public_key, ), ) rep_socket = context.socket(azmq.REP) try: req_socket.connect(ENDPOINT) rep_socket.bind(ENDPOINT) await req_socket.send_multipart([b'my', b'request']) with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(rep_socket.recv_multipart(), 0.25) finally: req_socket.close() rep_socket.close()
async def test_tcp_socket_plain_client( event_loop, socket_factory, connect_or_bind, ): rep_socket = socket_factory.create(zmq.REP) rep_socket.plain_server = True connect_or_bind(rep_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): frames = rep_socket.recv_multipart() assert frames == [b'my', b'question'] rep_socket.send_multipart([b'your', b'answer']) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket( socket_type=azmq.REQ, mechanism=PlainClient( username='******', password='******', ), ) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await asyncio.wait_for( socket.send_multipart([b'my', b'question']), 1, ) frames = await asyncio.wait_for(socket.recv_multipart(), 1) assert frames == [b'your', b'answer']
async def test_req_rep_multiple_receivers(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) rep_socket = context.socket(azmq.REP) rep_socket_2 = context.socket(azmq.REP) try: req_socket.bind(endpoint) rep_socket.connect(endpoint) rep_socket_2.connect(endpoint) multiplexer = Multiplexer() multiplexer.add_socket(rep_socket) multiplexer.add_socket(rep_socket_2) await fivesec(req_socket.send_multipart([b'my', b'request'])) (socket, message), = await fivesec(multiplexer.recv_multipart()) assert message == [b'my', b'request'] await fivesec(socket.send_multipart([b'my', b'response'])) message = await req_socket.recv_multipart() assert message == [b'my', b'response'] finally: req_socket.close() rep_socket.close() rep_socket_2.close()
async def test_pub_sub_subscribe_while_connected(event_loop, endpoint): async with azmq.Context() as context: pub_socket = context.socket(azmq.PUB) sub_socket = context.socket(azmq.SUB) await sub_socket.subscribe(b'h') try: pub_socket.bind(endpoint) sub_socket.connect(endpoint) # Let's start a task that sends messages on the pub socket. async def publish(): while True: await fivesec(pub_socket.send_multipart([b'hello'])) await fivesec(pub_socket.send_multipart([b'world'])) publish_task = asyncio.ensure_future(publish()) try: message = await fivesec(sub_socket.recv_multipart()) assert message == [b'hello'] await fivesec(sub_socket.unsubscribe(b'h')) await fivesec(sub_socket.subscribe(b'w')) while message == [b'hello']: message = await fivesec(sub_socket.recv_multipart()) assert message == [b'world'] finally: publish_task.cancel() finally: pub_socket.close() sub_socket.close()
async def test_pub_sub_spam_subscribe_after(event_loop, endpoint): async with azmq.Context() as context: pub_socket = context.socket(azmq.PUB) sub_sockets = [context.socket(azmq.SUB) for _ in range(10)] try: async def send(): pub_socket.bind(endpoint) while True: await fivesec(pub_socket.send_multipart([b'a', b'b'])) async def recv(socket): socket.connect(endpoint) await socket.subscribe(b'a') message = await fivesec(socket.recv_multipart()) assert message == [b'a', b'b'] send_task = asyncio.ensure_future(send()) recv_tasks = [ asyncio.ensure_future(recv(socket)) for socket in sub_sockets ] try: await fivesec(asyncio.gather(*recv_tasks)) finally: send_task.cancel() finally: pub_socket.close() for socket in sub_sockets: socket.close()
async def test_push_pull(event_loop, endpoint): async with azmq.Context() as context: push_socket = context.socket(azmq.PUSH) pull_socket_1 = context.socket(azmq.PULL) pull_socket_2 = context.socket(azmq.PULL) try: push_socket.bind(endpoint) pull_socket_1.connect(endpoint) pull_socket_2.connect(endpoint) multiplexer = Multiplexer() multiplexer.add_socket(pull_socket_1) multiplexer.add_socket(pull_socket_2) await fivesec(push_socket.send_multipart([b'a', b'1'])) await fivesec(push_socket.send_multipart([b'b', b'2'])) messages = [] while len(messages) < 2: results = await fivesec(multiplexer.recv_multipart()) messages.extend(tuple(x) for _, x in results) assert set(messages) == {(b'a', b'1'), (b'b', b'2')} finally: push_socket.close() pull_socket_1.close() pull_socket_2.close()
async def test_req_router_multiple_receivers(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) router_socket = context.socket(azmq.ROUTER) router_socket_2 = context.socket(azmq.ROUTER) router_sockets = {router_socket, router_socket_2} try: req_socket.identity = b'me' req_socket.bind(endpoint) router_socket.connect(endpoint) router_socket_2.connect(endpoint) multiplexer = Multiplexer() multiplexer.add_socket(router_socket_2) multiplexer.add_socket(router_socket) await fivesec(req_socket.send_multipart([b'my', b'request'])) (socket, message), = await fivesec(multiplexer.recv_multipart()) assert message == [b'me', b'', b'my', b'request'] other_socket = next(iter(router_sockets - {socket})) task = asyncio.ensure_future(req_socket.recv_multipart()) await fivesec( other_socket.send_multipart([b'me', b'', b'some', b'error'])) await fivesec( socket.send_multipart([b'me', b'', b'my', b'response'])) message = await fivesec(task) assert message == [b'my', b'response'] finally: req_socket.close() router_socket.close() router_socket_2.close()
async def test_req_rep_invalid_binds(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) rep_socket = context.socket(azmq.REP) try: # We trigger all kind of binding errors: invalid endpoints, # duplicate ports, and so on. req_socket.bind('ipc:///non-existing/path') req_socket.bind(endpoint) # This one is special and fails immediately. with pytest.raises(UnsupportedSchemeError): req_socket.bind('foo://some-endpoint') # This bind is valid and the test should work nonetheless. req_socket.bind(endpoint) rep_socket.connect(endpoint) await fivesec(req_socket.send_multipart([b'my', b'request'])) message = await fivesec(rep_socket.recv_multipart()) assert message == [b'my', b'request'] await fivesec(rep_socket.send_multipart([b'my', b'response'])) message = await req_socket.recv_multipart() assert message == [b'my', b'response'] finally: req_socket.close() rep_socket.close()
async def test_tcp_socket_curve_client( event_loop, socket_factory, connect_or_bind, ): rep_socket = socket_factory.create(zmq.REP) public, secret = curve_gen_keypair() rep_socket.curve_secretkey = secret rep_socket.curve_publickey = public rep_socket.curve_server = True connect_or_bind(rep_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): frames = rep_socket.recv_multipart() assert frames == [b'my', b'question'] rep_socket.send_multipart([b'your', b'answer']) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket( socket_type=azmq.REQ, mechanism=CurveClient(server_key=public), ) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await asyncio.wait_for( socket.send_multipart([b'my', b'question']), 1, ) frames = await asyncio.wait_for(socket.recv_multipart(), 1) assert frames == [b'your', b'answer']
async def test_plain_valid_password(event_loop): async with azmq.Context() as context: authenticator = ZAPAuthenticator(context) authenticator.add_user(username='******', password='******') context.set_zap_authenticator(authenticator) req_socket = context.socket( azmq.REQ, mechanism=PlainClient(username='******', password='******'), ) rep_socket = context.socket(azmq.REP, mechanism=PlainServer()) try: req_socket.connect(ENDPOINT) rep_socket.bind(ENDPOINT) await req_socket.send_multipart([b'my', b'request']) message = await onesec(rep_socket.recv_multipart()) assert message == [b'my', b'request'] await rep_socket.send_multipart([b'my', b'response']) message = await onesec(req_socket.recv_multipart()) assert message == [b'my', b'response'] finally: req_socket.close() rep_socket.close()
async def test_zap_successful_authentication_after_invalid_request(event_loop): class MyZAPAuthenticator(BaseZAPAuthenticator): async def on_request(self, *args, **kwargs): return 'bob', {b'foo': b'bar'} async with azmq.Context() as context: async with MyZAPAuthenticator(context=context): async with ZAPClient(context=context) as zap_client: async with context.socket(azmq.DEALER) as socket: socket.connect(ZAP_INPROC_ENDPOINT) await socket.send_multipart([b'invalid', b'data']) username, metadata = await asyncio.wait_for( zap_client.authenticate( domain='domain', address='127.0.0.1', identity=b'bob', mechanism=b'CURVE', credentials=[b'mycred', b'value'], ), 1, ) assert username == 'bob' assert metadata == {b'foo': b'bar'}
async def test_req_invalid_operation(event_loop, endpoint): async with azmq.Context() as context: async with context.socket(azmq.REQ) as req_socket: req_socket.connect(endpoint) await fivesec(req_socket.send_multipart([b'my', b'request'])) with pytest.raises(InvalidOperation): await fivesec(req_socket.send_multipart([b'my', b'request']))
async def test_tcp_pull_socket(event_loop, socket_factory, connect_or_bind): pull_socket = socket_factory.create(zmq.PUSH) connect_or_bind(pull_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): pull_socket.send_multipart([b'hello', b'world']) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.PULL) connect_or_bind(socket, 'tcp://127.0.0.1:3333') message = await socket.recv_multipart() assert message == [b'hello', b'world']
async def test_zap_unknown_request(event_loop): class MyZAPAuthenticator(BaseZAPAuthenticator): async def on_request(self, *args, **kwargs): return 'bob', {b'foo': b'bar'} async with azmq.Context() as context: async with ZAPClient(context=context) as zap_client: async with context.socket(azmq.ROUTER) as socket: socket.bind(ZAP_INPROC_ENDPOINT) task = asyncio.ensure_future( zap_client.authenticate( domain='domain', address='127.0.0.1', identity=b'bob', mechanism=b'CURVE', credentials=[b'mycred', b'value'], )) identity, *_ = await asyncio.wait_for( socket.recv_multipart(), 1, ) await socket.send_multipart([ identity, b'', b'1.0', b'my_request', b'200', b'all good', b'userid', b'', ]) assert not task.done() task.cancel() async with MyZAPAuthenticator(context=context): username, metadata = await asyncio.wait_for( zap_client.authenticate( domain='domain', address='127.0.0.1', identity=b'bob', mechanism=b'CURVE', credentials=[b'mycred', b'value'], ), 1, ) assert username == 'bob' assert metadata == {b'foo': b'bar'}
async def test_tcp_pub_socket(event_loop, socket_factory, connect_or_bind): sub_socket = socket_factory.create(zmq.SUB) sub_socket.setsockopt(zmq.SUBSCRIBE, b'a') connect_or_bind(sub_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): frames = sub_socket.recv_multipart() assert frames == [b'a', b'message'] with run_in_background(run) as thread_done_event: async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.PUB) connect_or_bind(socket, 'tcp://127.0.0.1:3333') while not thread_done_event.is_set(): await socket.send_multipart([b'a', b'message']) await socket.send_multipart([b'b', b'wrong'])
async def test_tcp_push_socket(event_loop, socket_factory, connect_or_bind): pull_socket = socket_factory.create(zmq.PULL) connect_or_bind(pull_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): assert pull_socket.poll(1000) == zmq.POLLIN message = pull_socket.recv_multipart() assert message == [b'hello', b'world'] with run_in_background(run) as event: async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.PUSH) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await socket.send_multipart([b'hello', b'world']) while not event.is_set(): await asyncio.sleep(0.1)
async def test_tcp_pair_socket(event_loop, socket_factory, connect_or_bind): pair_socket = socket_factory.create(zmq.PAIR) connect_or_bind(pair_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): assert pair_socket.poll(1000) == zmq.POLLIN message = pair_socket.recv_multipart() assert message == [b'hello', b'world'] pair_socket.send_multipart([b'my', b'message']) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.PAIR) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await socket.send_multipart([b'hello', b'world']) message = await asyncio.wait_for(socket.recv_multipart(), 1) assert message == [b'my', b'message']
async def test_xpub_xsub_unknown_subscription_message(event_loop, endpoint): async with azmq.Context() as context: xpub_socket = context.socket(azmq.XPUB) xsub_socket = context.socket(azmq.XSUB) try: xpub_socket.bind(endpoint) xsub_socket.connect(endpoint) # We force a protocol error by sending an unexpected message. await xsub_socket._fair_send([b'\2x']) await xsub_socket.send_multipart([b'\1a']) message = await fivesec(xpub_socket.recv_multipart()) assert message == [b'\1a'] finally: xpub_socket.close() xsub_socket.close()
async def test_push_pull_slow_connect(event_loop, endpoint): async with azmq.Context() as context: push_socket = context.socket(azmq.PUSH) pull_socket = context.socket(azmq.PULL) try: # The PUSH socket binds before the PULL sockets binds and sends # a message right away. It should block in this case. push_socket.bind(endpoint) await zerosec(push_socket.send_multipart([b'hello'])) # The PULL sockets finally binds, and should receive no message. pull_socket.connect(endpoint) await zerosec(pull_socket.recv_multipart()) finally: push_socket.close() pull_socket.close()
async def test_req_req_invalid_combination(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) req_socket_2 = context.socket(azmq.REQ) try: future = asyncio.Future(loop=event_loop) req_socket.on_connection_failure.connect(future.set_result) req_socket.connect(endpoint) req_socket_2.bind(endpoint) await fivesec(future) assert isinstance(future.result().error, ProtocolError) assert future.result().error.fatal finally: req_socket.close() req_socket_2.close()
async def test_tcp_big_messages(event_loop, socket_factory, connect_or_bind): rep_socket = socket_factory.create(zmq.REP) connect_or_bind(rep_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): frames = rep_socket.recv_multipart() assert frames == [b'1' * 500, b'2' * 100000] rep_socket.send_multipart([b'3' * 500, b'4' * 100000]) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.REQ) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await asyncio.wait_for( socket.send_multipart([b'1' * 500, b'2' * 100000]), 1, ) frames = await asyncio.wait_for(socket.recv_multipart(), 1) assert frames == [b'3' * 500, b'4' * 100000]
async def test_req_rep(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) rep_socket = context.socket(azmq.REP) try: req_socket.bind(endpoint) rep_socket.connect(endpoint) await fivesec(req_socket.send_multipart([b'my', b'request'])) message = await fivesec(rep_socket.recv_multipart()) assert message == [b'my', b'request'] await fivesec(rep_socket.send_multipart([b'my', b'response'])) message = await req_socket.recv_multipart() assert message == [b'my', b'response'] finally: req_socket.close() rep_socket.close()
async def test_req_router_invalid(event_loop, endpoint): async with azmq.Context() as context: req_socket = context.socket(azmq.REQ) router_socket = context.socket(azmq.ROUTER) try: req_socket.identity = b'bob' req_socket.bind(endpoint) router_socket.connect(endpoint) await fivesec(req_socket.send_multipart([b'my', b'request'])) message = await fivesec(router_socket.recv_multipart()) assert message == [b'bob', b'', b'my', b'request'] await fivesec(router_socket.send_multipart([b'bob', b'bug'])) await zerosec(req_socket.recv_multipart()) finally: req_socket.close() router_socket.close()
async def test_xpub_xsub_subscriptions(event_loop, endpoint): async with azmq.Context() as context: xpub_socket = context.socket(azmq.XPUB) xsub_socket = context.socket(azmq.XSUB) try: xpub_socket.bind(endpoint) xsub_socket.connect(endpoint) await xsub_socket.send_multipart([b'\1a']) message = await fivesec(xpub_socket.recv_multipart()) assert message == [b'\1a'] await xsub_socket.send_multipart([b'\0a']) message = await fivesec(xpub_socket.recv_multipart()) assert message == [b'\0a'] finally: xpub_socket.close() xsub_socket.close()
async def test_zap_default_authenticator_authentication_failure(event_loop): async with azmq.Context() as context: async with ZAPAuthenticator(context=context) as authenticator: authenticator.allow('192.168.0.1') async with ZAPClient(context=context) as zap_client: with pytest.raises(ZAPAuthenticationFailure) as error: await asyncio.wait_for( zap_client.authenticate( domain='domain', address='127.0.0.1', identity=b'bob', mechanism=b'CURVE', credentials=[b'mycred', b'value'], ), 1, ) assert error.value.code == 400
async def test_tcp_dealer_socket(event_loop, socket_factory, connect_or_bind): rep_socket = socket_factory.create(zmq.REP) connect_or_bind(rep_socket, 'tcp://127.0.0.1:3333', reverse=True) def run(): frames = rep_socket.recv_multipart() assert frames == [b'my', b'question'] rep_socket.send_multipart([b'your', b'answer']) with run_in_background(run): async with azmq.Context(loop=event_loop) as context: socket = context.socket(azmq.DEALER) connect_or_bind(socket, 'tcp://127.0.0.1:3333') await asyncio.wait_for( socket.send_multipart([b'', b'my', b'question']), 1, ) frames = await asyncio.wait_for(socket.recv_multipart(), 1) assert frames == [b'', b'your', b'answer']
async def test_zap_default_authenticator_authentication_success(event_loop): async with azmq.Context() as context: async with ZAPAuthenticator(context=context) as authenticator: authenticator.deny('192.168.0.1') async with ZAPClient(context=context) as zap_client: username, metadata = await asyncio.wait_for( zap_client.authenticate( domain='domain', address='127.0.0.1', identity=b'bob', mechanism=b'NULL', credentials=[], ), 1, ) assert username == '' assert metadata == {}
async def test_push_pull_slow_bind(event_loop, endpoint): async with azmq.Context() as context: push_socket = context.socket(azmq.PUSH) pull_socket = context.socket(azmq.PULL) try: # The PUSH socket connects before the PULL sockets binds and sends # a message right away. push_socket.connect(endpoint) await fivesec(push_socket.send_multipart([b'hello'])) # The PULL sockets finally binds, and should receive the message, # even late. pull_socket.bind(endpoint) message = await fivesec(pull_socket.recv_multipart()) assert message == [b'hello'] finally: push_socket.close() pull_socket.close()
async def test_ip_deny(event_loop): async with azmq.Context() as context: authenticator = ZAPAuthenticator(context) authenticator.deny('127.0.0.1') context.set_zap_authenticator(authenticator) req_socket = context.socket(azmq.REQ) rep_socket = context.socket(azmq.REP) try: req_socket.connect(ENDPOINT) rep_socket.bind(ENDPOINT) await req_socket.send_multipart([b'my', b'request']) with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(rep_socket.recv_multipart(), 0.25) finally: req_socket.close() rep_socket.close()