def test_consume_end_of_stream(self):
stream = ReceiveStream(1)
stream.close()
message_type = Mock()
assert list(stream.consume(message_type)) == []
def test_write_to_closed_stream(self):
stream = ReceiveStream(1)
assert stream.buffer.empty()
stream.close()
stream.write(b"\x00\x00\x00")
assert stream.buffer.empty()
def test_consume_uncompressed_message(self):
stream = ReceiveStream(1)
message_data = b"x"
message_type = Mock()
message = message_type()
stream.queue.put((False, message_data))
stream.close() # close stream so that consume exits
assert list(stream.consume(message_type)) == [message]
assert message.ParseFromString.call_args_list == [call(message_data)]