Function test_flight_do_put_limit

python/pyarrow/tests/test_flight.py:1882–1907

Try a simple do_put call with a size limit.

()

@pytest.mark.numpy
def test_flight_do_put_limit():
    """Try a simple do_put call with a size limit."""
    large_batch = pa.RecordBatch.from_arrays([
        pa.array(np.ones(768, dtype=np.int64())),
    ], names=['a'])

    with EchoFlightServer() as server, \
        FlightClient(('localhost', server.port),
                     write_size_limit_bytes=4096) as client:
        writer, metadata_reader = client.do_put(
            flight.FlightDescriptor.for_path(''),
            large_batch.schema)
        with writer:
            with pytest.raises(flight.FlightWriteSizeExceededError,
                               match="exceeded soft limit") as excinfo:
                writer.write_batch(large_batch)
            assert excinfo.value.limit == 4096
            smaller_batches = [
                large_batch.slice(0, 384),
                large_batch.slice(384),
            ]
            for batch in smaller_batches:
                writer.write_batch(batch)
        expected = pa.Table.from_batches([large_batch])
        actual = client.do_get(flight.Ticket(b'')).read_all()
        assert expected == actual


@pytest.mark.slow