Function test_flight_generator_stream_of_batches

python/pyarrow/tests/test_flight.py:1401–1414

Try downloading a flight of RecordBatches in a GeneratorStream.

()

def test_flight_generator_stream_of_batches():
    """Try downloading a flight of RecordBatches in a GeneratorStream."""
    data = pa.Table.from_arrays([
        pa.array(range(0, 10 * 1024))
    ], names=['a'])

    with EchoStreamFlightServer() as server, \
            FlightClient(('localhost', server.port)) as client:
        writer, _ = client.do_put(flight.FlightDescriptor.for_path('test'),
                                  data.schema)
        writer.write_table(data)
        writer.close()
        result = client.do_get(flight.Ticket(b'')).read_all()
        assert result.equals(data)


def test_flight_generator_stream_of_batches_with_dict():