pub/python-betterproto
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'client-streaming' into client-streaming-tests

Browse Source
This commit is contained in:
nat 2020-06-15 23:42:01 +02:00 committed by GitHub
commit 4e78fe9579
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 18 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
betterproto/grpc/util/async_channel.py
View File

@ -81,17 +81,10 @@ class AsyncChannel(AsyncIterable[T]):
""" """
def __init__( def __init__(
self, self, *, buffer_limit: int = 0, close: bool = False,
source: Union[Iterable[T], AsyncIterable[T]] = tuple(),
*,
buffer_limit: int = 0,
close: bool = False,
): ):
self._queue: asyncio.Queue[Union[T, object]] = asyncio.Queue(buffer_limit) self._queue: asyncio.Queue[Union[T, object]] = asyncio.Queue(buffer_limit)
self._closed = False self._closed = False
self._sending_task = (
asyncio.ensure_future(self.send_from(source, close)) if source else None
)
self._waiting_recievers: int = 0 self._waiting_recievers: int = 0
# Track whether flush has been invoked so it can only happen once # Track whether flush has been invoked so it can only happen once
self._flushed = False self._flushed = False
@ -132,7 +125,7 @@ class AsyncChannel(AsyncIterable[T]):
async def send_from( async def send_from(
self, source: Union[Iterable[T], AsyncIterable[T]], close: bool = False self, source: Union[Iterable[T], AsyncIterable[T]], close: bool = False
): ) -> "AsyncChannel[T]":
""" """
Iterates the given [Async]Iterable and sends all the resulting items. Iterates the given [Async]Iterable and sends all the resulting items.
If close is set to True then subsequent send calls will be rejected with a If close is set to True then subsequent send calls will be rejected with a
@ -153,8 +146,9 @@ class AsyncChannel(AsyncIterable[T]):
if close: if close:
# Complete the closing process # Complete the closing process
self.close() self.close()
return self
async def send(self, item: T): async def send(self, item: T) -> "AsyncChannel[T]":
""" """
Send a single item over this channel. Send a single item over this channel.
:param item: The item to send :param item: The item to send
@ -162,6 +156,7 @@ class AsyncChannel(AsyncIterable[T]):
if self._closed: if self._closed:
raise ChannelClosed("Cannot send through a closed channel") raise ChannelClosed("Cannot send through a closed channel")
await self._queue.put(item) await self._queue.put(item)
return self
async def recieve(self) -> Optional[T]: async def recieve(self) -> Optional[T]:
""" """
@ -185,8 +180,6 @@ class AsyncChannel(AsyncIterable[T]):
""" """
Close this channel to new items Close this channel to new items
""" """
if self._sending_task is not None:
self._sending_task.cancel()
self._closed = True self._closed = True
asyncio.ensure_future(self._flush_queue()) asyncio.ensure_future(self._flush_queue())

22
betterproto/tests/grpc/test_grpclib_client.py
View File

@ -1,3 +1,4 @@
import asyncio
from betterproto.tests.output_betterproto.service.service import ( from betterproto.tests.output_betterproto.service.service import (
DoThingResponse, DoThingResponse,
DoThingRequest, DoThingRequest,
@ -129,7 +130,10 @@ async def test_async_gen_for_stream_stream_request():
# Use an AsyncChannel to decouple sending and recieving, it'll send some_things # Use an AsyncChannel to decouple sending and recieving, it'll send some_things
# immediately and we'll use it to send more_things later, after recieving some # immediately and we'll use it to send more_things later, after recieving some
# results # results
request_chan = AsyncChannel(GetThingRequest(name) for name in some_things) request_chan = AsyncChannel()
send_initial_requests = asyncio.ensure_future(
request_chan.send_from(GetThingRequest(name) for name in some_things)
)
response_index = 0 response_index = 0
async for response in client.get_different_things(request_chan): async for response in client.get_different_things(request_chan):
assert response.name == expected_things[response_index] assert response.name == expected_things[response_index]
@ -138,13 +142,13 @@ async def test_async_gen_for_stream_stream_request():
if more_things: if more_things:
# Send some more requests as we recieve reponses to be sure coordination of # Send some more requests as we recieve reponses to be sure coordination of
# send/recieve events doesn't matter # send/recieve events doesn't matter
another_response = await request_chan.send( await request_chan.send(GetThingRequest(more_things.pop(0)))
GetThingRequest(more_things.pop(0)) elif not send_initial_requests.done():
) # Make sure the sending task it completed
if another_response is not None: await send_initial_requests
assert another_response.name == expected_things[response_index]
assert another_response.version == response_index
response_index += 1
else: else:
# No more things to send make sure channel is closed # No more things to send make sure channel is closed
await request_chan.close() request_chan.close()
assert response_index == len(
expected_things
), "Didn't recieve all exptected responses"