def test_load_rows_into_mem(synthetic_dataset, capacity, reader_factory):
    """
    See if inmem buffer length is minimum of capacity and input data length.
    See if keys are correct.
    Also see reader is stopped.
    """
    with reader_factory(synthetic_dataset.url,
                        schema_fields=[TestSchema.id,
                                       TestSchema.id_float,
                                       TestSchema.id_odd]) as reader:
        keys, buffer = _load_rows_into_mem(reader, torch.as_tensor, capacity)
        assert len(buffer[0]) == min(capacity, len(synthetic_dataset.data))
        assert list(keys)[0] == ("id")
        assert list(keys)[1] == ("id_float")
        assert list(keys)[2] == ("id_odd")
        assert reader.stopped

    See if inmem buffer length is minimum of capacity and input data length.
    See if keys are correct.
    Also see reader is stopped.
    
petastorm/tests/test_pytorch_dataloader.py