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()