Function test_doexchange_echo

python/pyarrow/tests/test_flight.py:2249–2289

Try a DoExchange echo server.

()

def test_doexchange_echo():
    """Try a DoExchange echo server."""
    data = pa.Table.from_arrays([
        pa.array(range(0, 10 * 1024))
    ], names=["a"])
    batches = data.to_batches(max_chunksize=512)

    with ExchangeFlightServer() as server, \
            FlightClient(("localhost", server.port)) as client:
        descriptor = flight.FlightDescriptor.for_command(b"echo")
        writer, reader = client.do_exchange(descriptor)
        with writer:
            # Read/write metadata before starting data.
            for i in range(10):
                buf = str(i).encode("utf-8")
                writer.write_metadata(buf)
                chunk = reader.read_chunk()
                assert chunk.data is None
                assert chunk.app_metadata == buf

            # Now write data without metadata.
            writer.begin(data.schema)
            num_batches = 0
            for batch in batches:
                writer.write_batch(batch)
                assert reader.schema == data.schema
                chunk = reader.read_chunk()
                assert chunk.data == batch
                assert chunk.app_metadata is None
                num_batches += 1
                assert reader.stats.num_record_batches == num_batches

            # And write data with metadata.
            for i, batch in enumerate(batches):
                buf = str(i).encode("utf-8")
                writer.write_with_metadata(batch, buf)
                chunk = reader.read_chunk()
                assert chunk.data == batch
                assert chunk.app_metadata == buf
                num_batches += 1
                assert reader.stats.num_record_batches == num_batches


def test_doexchange_echo_v4():