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_rx_support_request_channel_server_take_only_n(pipe: Tuple[RSocketServer, RSocketClient],
take_only_n):
server, client = pipe
received_messages = []
items_generated = 0
maximum_message_count = 3
wait_for_server_finish = asyncio.Event()
class Handler(BaseRequestHandler, DefaultSubscriber):
def on_next(self, value: Payload, is_complete=False):
received_messages.append(value)
if len(received_messages) < take_only_n:
self.subscription.request(1)
else:
self.subscription.cancel()
wait_for_server_finish.set()
def on_complete(self):
wait_for_server_finish.set()
async def on_error(self, error_code: ErrorCode, payload: Payload):
wait_for_server_finish.set()
def on_subscribe(self, subscription: Subscription):
super().on_subscribe(subscription)
subscription.request(1)
async def request_channel(self, payload: Payload) -> Tuple[Optional[Publisher], Optional[Subscriber]]:
return None, self
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
def generator():
nonlocal items_generated
for x in range(maximum_message_count):
items_generated += 1
yield Payload('Feed Item: {}'.format(x).encode('utf-8'))
await rx_client.request_channel(
Payload(b'request text'),
observable=rx.from_iterable(generator())
).pipe(
operators.to_list()
)
await wait_for_server_finish.wait()
maximum_message_received = min(maximum_message_count, take_only_n)
# assert items_generated == maximum_message_received # todo: Stop async generator on cancel from server requester
assert len(received_messages) == maximum_message_received
for i in range(maximum_message_received):
assert received_messages[i].data == ('Feed Item: %d' % i).encode()
async def test_rx_support_request_response_cancel_on_timeout(
pipe: Tuple[RSocketServer, RSocketClient]):
server, client = pipe
response_sent = False
class Handler(BaseRequestHandler):
async def request_response(self, payload: Payload) -> Future:
nonlocal response_sent
await asyncio.sleep(3)
response_sent = True
return create_future(Payload(b'response'))
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
with pytest.raises(Exception):
await asyncio.wait_for(
rx_client.request_response(Payload(b'request text')), 2)
assert not response_sent
async def test_rx_support_request_channel_with_error_from_requester(
pipe: Tuple[RSocketServer, RSocketClient], success_count,
request_limit):
server, client = pipe
responder_received_error = asyncio.Event()
server_received_messages = []
received_error = None
class ResponderSubscriber(DefaultSubscriber):
def on_subscribe(self, subscription: Subscription):
super().on_subscribe(subscription)
self.subscription.request(1)
def on_next(self, value, is_complete=False):
if len(value.data) > 0:
server_received_messages.append(value.data)
self.subscription.request(1)
def on_error(self, exception: Exception):
nonlocal received_error
received_error = exception
responder_received_error.set()
async def generator() -> AsyncGenerator[Tuple[Payload, bool], None]:
for x in range(success_count):
yield Payload('Feed Item: {}'.format(x).encode(
'utf-8')), x == success_count - 1
class Handler(BaseRequestHandler):
async def request_channel(
self, payload: Payload
) -> Tuple[Optional[Publisher], Optional[Subscriber]]:
return StreamFromAsyncGenerator(generator), ResponderSubscriber()
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
def test_observable(observer: Observer, scheduler: Optional[Scheduler]):
observer.on_error(Exception('Some error'))
return Disposable()
await rx_client.request_channel(
Payload(b'request text'),
observable=rx.create(test_observable),
request_limit=request_limit).pipe(
operators.map(lambda payload: payload.data), operators.to_list())
await responder_received_error.wait()
assert str(received_error) == 'Some error'
async def test_rx_support_request_stream_cancel_on_timeout(
pipe: Tuple[RSocketServer, RSocketClient]):
server, client = pipe
cancel_done = asyncio.Event()
stream_messages_sent_count = 0
class Handler(BaseRequestHandler, DefaultPublisherSubscription):
async def delayed_stream(self):
nonlocal stream_messages_sent_count
try:
await asyncio.sleep(3)
self._subscriber.on_next(Payload(b'success'))
stream_messages_sent_count += 1
except asyncio.CancelledError:
cancel_done.set()
def cancel(self):
self._task.cancel()
def request(self, n: int):
self._task = asyncio.create_task(self.delayed_stream())
async def request_stream(self, payload: Payload) -> Publisher:
return self
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
with pytest.raises(Exception):
await asyncio.wait_for(
rx_client.request_stream(Payload(b'request text')).pipe(
operators.to_list()), 2)
await cancel_done.wait()
assert stream_messages_sent_count == 0
async def test_rx_support_request_channel_response_take_only_n(pipe: Tuple[RSocketServer, RSocketClient],
take_only_n):
server, client = pipe
maximum_message_count = 4
wait_for_server_finish = asyncio.Event()
items_generated = 0
async def generator() -> AsyncGenerator[Tuple[Payload, bool], None]:
nonlocal items_generated
for x in range(maximum_message_count):
items_generated += 1
yield Payload('Feed Item: {}'.format(x).encode('utf-8')), x == maximum_message_count - 1
class Handler(BaseRequestHandler):
async def request_channel(self, payload: Payload) -> Tuple[Optional[Publisher], Optional[Subscriber]]:
def set_server_finished(): wait_for_server_finish.set()
return StreamFromAsyncGenerator(generator,
on_cancel=set_server_finished,
on_complete=set_server_finished), None
server.set_handler_using_factory(Handler)
rx_client = RxRSocket(client)
received_messages = await rx_client.request_channel(
Payload(b'request text'),
request_limit=1
).pipe(
operators.map(lambda payload: payload.data),
operators.take(take_only_n),
operators.to_list()
)
if take_only_n > 0:
await wait_for_server_finish.wait()
maximum_message_received = min(maximum_message_count, take_only_n)
assert len(received_messages) == maximum_message_received, 'Received message count wrong'
assert items_generated == maximum_message_received, 'Received message count wrong'
for i in range(maximum_message_received):
assert received_messages[i] == ('Feed Item: %d' % i).encode()
async def main(server_port):
logging.info('Connecting to server at localhost:%s', server_port)
connection = await asyncio.open_connection('localhost', server_port)
async with RSocketClient(single_transport_provider(
TransportTCP(*connection)),
metadata_encoding=WellKnownMimeTypes.
MESSAGE_RSOCKET_COMPOSITE_METADATA) as client:
rx_client = RxRSocket(client)
await request_response(rx_client)
await request_stream(rx_client)
await request_slow_stream(rx_client)
await request_channel(rx_client)
await request_stream_invalid_login(rx_client)
await request_fragmented_stream(rx_client)
await metadata_push(rx_client, b'audit info')
await request_last_metadata(rx_client)
await fire_and_forget(rx_client, b'aux data')
await request_last_fnf(rx_client)