async def test_connection_failure_during_stream(unused_tcp_port, generate_test_certificates,
transport_id, start_service, start_client):
logging.info('Testing transport %s on port %s', transport_id, unused_tcp_port)
server_container = ServerContainer()
wait_for_server = Event()
service_closer = await start_service(wait_for_server, server_container, unused_tcp_port, generate_test_certificates)
client = await start_client(unused_tcp_port, generate_test_certificates)
try:
async with AwaitableRSocket(client) as async_client:
await wait_for_server.wait()
wait_for_server.clear()
with pytest.raises(RSocketProtocolError) as exc_info:
await asyncio.gather(
async_client.request_stream(Payload(b'request 1')),
force_closing_connection(server_container.transport, timedelta(seconds=2)))
assert exc_info.value.data == 'Connection error'
assert exc_info.value.error_code == ErrorCode.CONNECTION_ERROR
await server_container.server.close() # cleanup async tasks from previous server to avoid errors (?)
await wait_for_server.wait()
response2 = await async_client.request_response(Payload(b'request 2'))
assert response2.data == b'data: request 2 server 2'
finally:
await server_container.server.close()
await service_closer()
async def test_send_frame_for_non_existing_stream(pipe_tcp, caplog):
(client, server) = pipe_tcp
done = asyncio.Event()
class Handler(BaseRequestHandler):
async def request_fire_and_forget(self, payload: Payload):
done.set()
async def request_response(self,
payload: Payload) -> Awaitable[Payload]:
return create_future(Payload(b'response'))
server.set_handler_using_factory(Handler)
bad_client = MisbehavingRSocket(client._transport)
client.fire_and_forget(Payload())
await bad_client.send_frame(to_payload_frame(145, Payload()))
await client.request_response(Payload(b'request'))
await done.wait()
records = caplog.get_records('call')
dropped_frame_log = [
record for record in records
if 'Dropping frame from unknown stream 145' in record.message
]
assert len(dropped_frame_log) > 0
async def test_request_response_with_client_and_server_side_lease_works(
lazy_pipe):
class Handler(BaseRequestHandler):
async def request_response(self, request: Payload):
return future_from_payload(request)
async with PeriodicalLeasePublisher(
maximum_request_count=2,
maximum_lease_time=timedelta(seconds=3),
wait_between_leases=timedelta(seconds=2)) as client_leases:
async with PeriodicalLeasePublisher(
maximum_request_count=2,
maximum_lease_time=timedelta(seconds=3),
wait_between_leases=timedelta(seconds=2)) as server_leases:
async with lazy_pipe(client_arguments={
'honor_lease': True,
'handler_factory': Handler,
'lease_publisher': client_leases
},
server_arguments={
'honor_lease': True,
'handler_factory': Handler,
'lease_publisher': server_leases
}) as (server, client):
for x in range(3):
response = await client.request_response(
Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
for x in range(3):
response = await server.request_response(
Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
async def test_authentication_success_on_setup(lazy_pipe):
class Handler(BaseRequestHandler):
def __init__(self, socket):
super().__init__(socket)
self._authenticated = False
async def on_setup(self, data_encoding: bytes,
metadata_encoding: bytes, payload: Payload):
composite_metadata = self._parse_composite_metadata(
payload.metadata)
authentication: AuthenticationSimple = composite_metadata.items[
0].authentication
if authentication.username != b'user' or authentication.password != b'12345':
raise Exception('Authentication rejected')
self._authenticated = True
async def request_response(self,
payload: Payload) -> Awaitable[Payload]:
if not self._authenticated:
raise RSocketApplicationError("Not authenticated")
return create_future(Payload(b'response'))
async with lazy_pipe(client_arguments={
'setup_payload':
Payload(metadata=composite(authenticate_simple('user', '12345')))
},
server_arguments={'handler_factory':
Handler}) as (server, client):
result = await client.request_response(Payload(b'request'))
assert result.data == b'response'
async def test_request_response_awaitable_wrapper(pipe):
class Handler(BaseRequestHandler):
async def request_response(self, request: Payload):
return future_from_payload(request)
server, client = pipe
server._handler = Handler(server)
response = await AwaitableRSocket(client).request_response(Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
async def _receiver(self):
try:
connection = Connection()
while True:
data = await self._reader.read(1024)
if not data:
self._writer.close()
break
frames = connection.receive_data(data)
for frame in frames:
stream = frame.stream_id
if stream and stream in self._streams:
self._streams[stream].frame_received(frame)
continue
if isinstance(frame, CancelFrame):
pass
elif isinstance(frame, ErrorFrame):
pass
elif isinstance(frame, KeepAliveFrame):
frame.flags_respond = False
self.send_frame(frame)
elif isinstance(frame, LeaseFrame):
pass
elif isinstance(frame, MetadataPushFrame):
pass
elif isinstance(frame, RequestChannelFrame):
pass
elif isinstance(frame, RequestFireAndForgetFrame):
pass
elif isinstance(frame, RequestResponseFrame):
stream = frame.stream_id
self._streams[stream] = RequestResponseResponder(
stream, self,
self._handler.request_response(
Payload(frame.data, frame.metadata)))
elif isinstance(frame, RequestStreamFrame):
stream = frame.stream_id
self._streams[stream] = RequestStreamResponder(
stream, self,
self._handler.request_stream(
Payload(frame.data, frame.metadata)))
elif isinstance(frame, RequestSubscriptionFrame):
pass
elif isinstance(frame, RequestNFrame):
pass
elif isinstance(frame, ResponseFrame):
pass
elif isinstance(frame, SetupFrame):
if frame.flags_lease:
lease = LeaseFrame()
lease.time_to_live = 10000
lease.number_of_requests = 100
self.send_frame(lease)
except asyncio.CancelledError:
pass
async def test_request_response_repeated(pipe):
class Handler(BaseRequestHandler):
async def request_response(self, request: Payload):
return future_from_payload(request)
server, client = pipe
server._handler = Handler(server)
for x in range(2):
response = await client.request_response(Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
async def test_authentication_failure_on_setup(lazy_pipe):
received_error_event = Event()
received_error: Optional[tuple] = None
class ServerHandler(BaseRequestHandler):
def __init__(self, socket):
super().__init__(socket)
self._authenticated = False
async def on_setup(self, data_encoding: bytes,
metadata_encoding: bytes, payload: Payload):
composite_metadata = self._parse_composite_metadata(
payload.metadata)
authentication: AuthenticationSimple = composite_metadata.items[
0].authentication
if authentication.username != b'user' or authentication.password != b'12345':
raise Exception('Authentication error')
self._authenticated = True
async def request_response(self,
payload: Payload) -> Awaitable[Payload]:
if not self._authenticated:
raise Exception("Not authenticated")
future = asyncio.get_event_loop().create_future()
future.set_result(Payload(b'response'))
return future
class ClientHandler(BaseRequestHandler):
async def on_error(self, error_code: ErrorCode, payload: Payload):
nonlocal received_error
received_error = (error_code, payload)
received_error_event.set()
async with lazy_pipe(client_arguments={
'handler_factory':
ClientHandler,
'setup_payload':
Payload(metadata=composite(
authenticate_simple('user', 'wrong_password')))
},
server_arguments={'handler_factory':
ServerHandler}) as (server, client):
with pytest.raises(RuntimeError):
await client.request_response(Payload(b'request'))
await received_error_event.wait()
assert received_error[0] == ErrorCode.REJECTED_SETUP
assert received_error[1] == Payload(b'Authentication error', b'')
async def test_request_response_repeated(pipe):
class Handler(BaseRequestHandler):
def request_response(self, request: Payload):
future = asyncio.Future()
future.set_result(Payload(b'data: ' + request.data,
b'meta: ' + request.metadata))
return future
server, client = pipe
server._handler = Handler(server)
for x in range(2):
response = await client.request_response(Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
async def test_request_stream_properly_finished(pipe: Tuple[RSocketServer,
RSocketClient],
complete_inline):
server, client = pipe
class Handler(BaseRequestHandler):
async def request_stream(self, payload: Payload) -> Publisher:
return StreamFromAsyncGenerator(self.feed)
async def feed(self):
for x in range(3):
value = Payload('Feed Item: {}'.format(x).encode('utf-8'))
yield value, complete_inline and x == 2
if not complete_inline:
yield None, True
server.set_handler_using_factory(Handler)
result = await AwaitableRSocket(client).request_stream(Payload())
assert len(result) == 3
assert result[0].data == b'Feed Item: 0'
assert result[1].data == b'Feed Item: 1'
assert result[2].data == b'Feed Item: 2'
async def test_fragmented_stream(pipe: Tuple[RSocketServer, RSocketClient]):
server, client = pipe
fragments_sent = 0
def generator() -> Generator[Tuple[Payload, bool], None, None]:
for i in range(3):
yield Payload(
ensure_bytes('some long data which should be fragmented %s' %
i)), i == 2
class StreamFragmentedCounter(StreamFromGenerator):
def _send_to_subscriber(self, payload: Payload, is_complete=False):
nonlocal fragments_sent
fragments_sent += 1
return super()._send_to_subscriber(payload, is_complete)
class Handler(BaseRequestHandler):
async def request_stream(self, payload: Payload) -> Publisher:
return StreamFragmentedCounter(generator, fragment_size=6)
server.set_handler_using_factory(Handler)
received_messages = await AwaitableRSocket(client).request_stream(
Payload())
assert len(received_messages) == 3
assert received_messages[
0].data == b'some long data which should be fragmented 0'
assert received_messages[
1].data == b'some long data which should be fragmented 1'
assert received_messages[
2].data == b'some long data which should be fragmented 2'
assert fragments_sent == 24
async def test_request_stream_immediately_completed_by_server_without_payloads(
pipe: Tuple[RSocketServer, RSocketClient]):
server, client = pipe
stream_done = asyncio.Event()
class Handler(BaseRequestHandler, DefaultPublisherSubscription):
def request(self, n: int):
self._subscriber.on_complete()
async def request_stream(self, payload: Payload) -> Publisher:
return self
class StreamSubscriber(DefaultSubscriber):
def __init__(self):
super().__init__()
self.received_messages: List[Payload] = []
def on_next(self, value, is_complete=False):
self.received_messages.append(value)
self.subscription.cancel()
logging.info(value)
def on_complete(self):
stream_done.set()
server.set_handler_using_factory(Handler)
stream_subscriber = StreamSubscriber()
client.request_stream(Payload()).subscribe(stream_subscriber)
await stream_done.wait()
assert len(stream_subscriber.received_messages) == 0
def frame_received(self, frame):
if isinstance(frame, ResponseFrame):
self.set_result(Payload(frame.data, frame.metadata))
self.socket.finish_stream(self.stream)
elif isinstance(frame, ErrorFrame):
self.set_exception(RuntimeError(frame.data))
self.socket.finish_stream(self.stream)
async def test_request_response_cancellation(pipe):
server_future = create_future()
class Handler(BaseRequestHandler):
async def request_response(self, payload: Payload):
# return a future that will never complete.
return server_future
server, client = pipe
server._handler = Handler(server)
future = client.request_response(Payload())
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(asyncio.shield(server_future), 0.1)
assert not server_future.cancelled()
future.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(asyncio.shield(server_future), 0.1)
with pytest.raises(asyncio.CancelledError):
await future
async def test_valid_authentication_in_routing_handler(lazy_pipe):
router = RequestRouter()
async def authenticate(path: str, authentication: Authentication):
if not isinstance(
authentication,
AuthenticationSimple) or authentication.password != b'pass':
raise Exception('Invalid credentials')
@router.response('test.path')
async def response():
return create_future(Payload(b'result'))
def handler_factory(socket):
return RoutingRequestHandler(socket,
router,
authentication_verifier=authenticate)
async with lazy_pipe(client_arguments={
'metadata_encoding':
WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA
},
server_arguments={'handler_factory': handler_factory
}) as (server, client):
result = await RxRSocket(client).request_response(
Payload(metadata=composite(route('test.path'),
authenticate_simple('user', 'pass'))))
assert result.data == b'result'
async def test_routed_request_channel_properly_finished(lazy_pipe):
router = RequestRouter()
def handler_factory(socket):
return RoutingRequestHandler(socket, router)
def feed():
for x in range(3):
yield Payload('Feed Item: {}'.format(x).encode('utf-8')), x == 2
@router.channel('test.path')
async def response_stream():
return StreamFromGenerator(feed), DefaultSubscriber()
async with lazy_pipe(client_arguments={
'metadata_encoding':
WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA
},
server_arguments={'handler_factory': handler_factory
}) as (server, client):
received_messages = await AwaitableRSocket(client).request_channel(
Payload(metadata=composite(route('test.path'))))
assert len(received_messages) == 3
assert received_messages[0].data == b'Feed Item: 0'
assert received_messages[1].data == b'Feed Item: 1'
assert received_messages[2].data == b'Feed Item: 2'
async def test_routed_request_response_properly_finished_accept_payload_and_metadata(
lazy_pipe):
router = RequestRouter()
def handler_factory(socket):
return RoutingRequestHandler(socket, router)
@router.response('test.path')
async def response(payload: Payload,
composite_metadata: CompositeMetadata):
return create_future(
Payload(('Response %s' % payload.data.decode()).encode(),
composite_metadata.items[0].tags[0]))
async with lazy_pipe(client_arguments={
'metadata_encoding':
WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA
},
server_arguments={'handler_factory': handler_factory
}) as (server, client):
result = await client.request_response(
Payload(data=b'request', metadata=composite(route('test.path'))))
assert result.data == b'Response request'
assert result.metadata == b'test.path'
async def feed(self):
for x in range(3):
value = Payload('Feed Item: {}'.format(x).encode('utf-8'))
yield value, complete_inline and x == 2
if not complete_inline:
yield None, True
async def test_request_response_bidirectional(pipe):
class ServerHandler(BaseRequestHandler):
@staticmethod
def future_done(other: asyncio.Future, current: asyncio.Future):
if current.cancelled():
other.set_exception(RuntimeError('Canceled.'))
elif current.exception():
other.set_exception(current.exception())
else:
payload = current.result()
payload.data = b'(server ' + payload.data + b')'
payload.metadata = b'(server ' + payload.metadata + b')'
other.set_result(payload)
async def request_response(self, payload: Payload):
future = create_future()
self.socket.request_response(payload).add_done_callback(
functools.partial(self.future_done, future))
return future
class ClientHandler(BaseRequestHandler):
async def request_response(self, payload: Payload):
return create_future(Payload(b'(client ' + payload.data + b')',
b'(client ' + payload.metadata + b')'))
server, client = pipe
server._handler = ServerHandler(server)
client._handler = ClientHandler(client)
response = await client.request_response(Payload(b'data', b'metadata'))
assert response.data == b'(server (client data))'
assert response.metadata == b'(server (client metadata))'
async def main(server_port):
logging.info('Connecting to server at localhost:%s', server_port)
client_configuration = QuicConfiguration(
is_client=True
)
ca_file_path = Path(__file__).parent / 'certificates' / 'pycacert.pem'
client_configuration.load_verify_locations(cafile=str(ca_file_path))
async with rsocket_connect('localhost', server_port,
configuration=client_configuration) as transport:
async with RSocketClient(single_transport_provider(transport)) as client:
payload = Payload(b'%Y-%m-%d %H:%M:%S')
async def run_request_response():
try:
while True:
result = await client.request_response(payload)
logging.info('Response: {}'.format(result.data))
await asyncio.sleep(1)
except asyncio.CancelledError:
pass
task = asyncio.create_task(run_request_response())
await asyncio.sleep(5)
task.cancel()
await task
async def test_request_stream_and_cancel_after_first_message(
pipe: Tuple[RSocketServer, RSocketClient]):
server, client = pipe
stream_canceled = asyncio.Event()
async def feed():
for x in range(3):
yield Payload('Feed Item: {}'.format(x).encode('utf-8')), x == 2
await asyncio.sleep(1)
class Handler(BaseRequestHandler):
async def request_stream(self, payload: Payload) -> Publisher:
return StreamFromAsyncGenerator(
feed, on_cancel=lambda: stream_canceled.set())
class StreamSubscriber(DefaultSubscriber):
def __init__(self):
super().__init__()
self.received_messages: List[Payload] = []
def on_next(self, value, is_complete=False):
self.received_messages.append(value)
self.subscription.cancel()
logging.info(value)
server.set_handler_using_factory(Handler)
stream_subscriber = StreamSubscriber()
client.request_stream(Payload()).subscribe(stream_subscriber)
await stream_canceled.wait()
assert len(stream_subscriber.received_messages) == 1
assert stream_subscriber.received_messages[0].data == b'Feed Item: 0'
async def request_fragmented_stream(client: RxRSocket):
payload = Payload(
b'The quick brown fox',
composite(route('fragmented_stream'),
authenticate_simple('user', '12345')))
result = await client.request_stream(payload).pipe(operators.to_list())
print(result)
async def test_routed_fire_and_forget(lazy_pipe):
router = RequestRouter()
received_data = None
received = asyncio.Event()
def handler_factory(socket):
return RoutingRequestHandler(socket, router)
@router.fire_and_forget('test.path')
async def fire_and_forget(payload):
nonlocal received_data
received_data = payload.data
received.set()
async with lazy_pipe(client_arguments={
'metadata_encoding':
WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA
},
server_arguments={'handler_factory': handler_factory
}) as (server, client):
client.fire_and_forget(
Payload(b'request data', composite(route('test.path'))))
await received.wait()
assert received_data == b'request data'
async def request_last_metadata(client: RxRSocket):
payload = Payload(metadata=composite(route('last_metadata_push'),
authenticate_simple('user', '12345')))
result = await client.request_response(payload).pipe()
assert result.data == b'audit info'
async def test_invalid_authentication_in_routing_handler(lazy_pipe):
router = RequestRouter()
async def authenticate(path: str, authentication: Authentication):
if not isinstance(
authentication,
AuthenticationSimple) or authentication.password != b'pass':
raise Exception('Invalid credentials')
@router.channel('test.path')
async def request_channel():
raise Exception('error from server')
def handler_factory(socket):
return RoutingRequestHandler(socket,
router,
authentication_verifier=authenticate)
async with lazy_pipe(client_arguments={
'metadata_encoding':
WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA
},
server_arguments={'handler_factory': handler_factory
}) as (server, client):
with pytest.raises(Exception) as exc_info:
await RxRSocket(client).request_channel(
Payload(metadata=composite(
route('test.path'),
authenticate_simple('user', 'wrong_password'))))
assert str(exc_info.value) == 'Invalid credentials'
async def test_rx_support_request_stream_with_error(pipe: Tuple[RSocketServer,
RSocketClient],
success_count,
request_limit):
server, client = pipe
async def generator() -> AsyncGenerator[Tuple[Payload, bool], None]:
for x in range(success_count):
yield Payload('Feed Item: {}'.format(x).encode('utf-8')), False
raise Exception('Some error from responder')
class Handler(BaseRequestHandler):
async def request_stream(self, payload: Payload) -> Publisher:
return StreamFromAsyncGenerator(generator)
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
with pytest.raises(Exception):
await rx_client.request_stream(
Payload(b'request text'), request_limit=request_limit).pipe(
operators.map(lambda payload: payload.data),
operators.to_list())
async def test_load_balancer_round_robin_request_response(
unused_tcp_port_factory):
clients = []
server_count = 3
request_count = 7
async with AsyncExitStack() as stack:
for i in range(server_count):
tcp_port = unused_tcp_port_factory()
_, client = await stack.enter_async_context(
pipe_factory_tcp(tcp_port,
server_arguments={
'handler_factory':
IdentifiedHandlerFactory(i,
Handler).factory
},
auto_connect_client=False))
clients.append(client)
round_robin = LoadBalancerRoundRobin(clients)
async with LoadBalancerRSocket(round_robin) as load_balancer_client:
results = await asyncio.gather(*[
load_balancer_client.request_response(
Payload(('request %d' % j).encode()))
for j in range(request_count)
])
assert results[0].data == b'data: request 0 server 0'
assert results[1].data == b'data: request 1 server 1'
assert results[2].data == b'data: request 2 server 2'
assert results[3].data == b'data: request 3 server 0'
assert results[4].data == b'data: request 4 server 1'
assert results[5].data == b'data: request 5 server 2'
assert results[6].data == b'data: request 6 server 0'
async def test_request_response_with_server_side_lease_works(lazy_pipe):
class Handler(BaseRequestHandler):
async def request_response(self, request: Payload):
return future_from_payload(request)
async with lazy_pipe(
client_arguments={'honor_lease': True},
server_arguments={
'handler_factory':
Handler,
'lease_publisher':
SingleLeasePublisher(maximum_lease_time=timedelta(seconds=3))
}) as (server, client):
for x in range(2):
response = await client.request_response(Payload(b'dog', b'cat'))
assert response == Payload(b'data: dog', b'meta: cat')
async def request_response(client: RSocketClient):
payload = Payload(
b'The quick brown fox',
composite(route('single_request'),
authenticate_simple('user', '12345')))
await client.request_response(payload)
async def test_load_balancer_random_request_response(unused_tcp_port_factory):
clients = []
server_count = 3
request_count = 70
async with AsyncExitStack() as stack:
for i in range(server_count):
tcp_port = unused_tcp_port_factory()
_, client = await stack.enter_async_context(
pipe_factory_tcp(tcp_port,
server_arguments={
'handler_factory':
IdentifiedHandlerFactory(i,
Handler).factory
},
auto_connect_client=False))
clients.append(client)
strategy = LoadBalancerRandom(clients)
async with LoadBalancerRSocket(strategy) as load_balancer_client:
results = await asyncio.gather(*[
load_balancer_client.request_response(
Payload(('request %d' % j).encode()))
for j in range(request_count)
])
server_ids = [payload.data.decode()[-1] for payload in results]
assert '0' in server_ids
assert '1' in server_ids
assert '2' in server_ids