def create_and_open():
underlying = self._underlying_factory(
subscription, fixed_partitions,
per_partition_flow_control_settings)
subscriber = SubscriberImpl(underlying, callback, self._executor)
future = StreamingPullFuture(subscriber)
subscriber.__enter__()
return future
def test_failed(subscriber: SubscriberImpl, async_subscriber, close_callback):
error = FailedPrecondition("bad read")
async_subscriber.read.side_effect = error
close_called = concurrent.futures.Future()
close_callback.side_effect = lambda manager, err: close_called.set_result(
None)
subscriber.add_close_callback(close_callback)
subscriber.__enter__()
async_subscriber.__aenter__.assert_called_once()
close_called.result()
async_subscriber.__aexit__.assert_called_once()
close_callback.assert_called_once_with(subscriber, error)
def subscribe(
self,
subscription: Union[SubscriptionPath, str],
callback: MessageCallback,
per_partition_flow_control_settings: FlowControlSettings,
fixed_partitions: Optional[Set[Partition]] = None,
) -> StreamingPullFuture:
if isinstance(subscription, str):
subscription = SubscriptionPath.parse(subscription)
underlying = self._underlying_factory(
subscription, fixed_partitions,
per_partition_flow_control_settings)
subscriber = SubscriberImpl(underlying, callback, self._executor)
future = StreamingPullFuture(subscriber)
subscriber.__enter__()
future.add_done_callback(lambda fut: self._try_remove_client(future))
return future
def test_messages_received(subscriber: SubscriberImpl, async_subscriber,
message_callback, close_callback):
message1 = Message(PubsubMessage(message_id="1")._pb, "", 0, None)
message2 = Message(PubsubMessage(message_id="2")._pb, "", 0, None)
counter = Box[int]()
counter.val = 0
async def on_read() -> Message:
counter.val += 1
if counter.val == 1:
return message1
if counter.val == 2:
return message2
await sleep_forever()
async_subscriber.read.side_effect = on_read
results = Queue()
message_callback.side_effect = lambda m: results.put(m.message_id)
subscriber.add_close_callback(close_callback)
subscriber.__enter__()
assert results.get() == "1"
assert results.get() == "2"
subscriber.close()
def test_init(subscriber: SubscriberImpl, async_subscriber, close_callback):
async_subscriber.read.side_effect = sleep_forever
subscriber.add_close_callback(close_callback)
subscriber.__enter__()
async_subscriber.__aenter__.assert_called_once()
subscriber.close()
async_subscriber.__aexit__.assert_called_once()
close_callback.assert_called_once_with(subscriber, None)
def subscriber(async_subscriber, message_callback, close_callback):
return SubscriberImpl(async_subscriber, message_callback,
ThreadPoolExecutor(max_workers=1))