def test_random_shuffling_buffer_can_add_retrieve_flags(buffer_type):
    """Check can_add/can_retrieve flags at all possible states"""
    q = buffer_type(shuffling_buffer_capacity=5, min_after_retrieve=3)

    # Empty buffer. Can start adding, nothing to retrieve yet
    assert q.size == 0
    assert q.can_add()
    assert not q.can_retrieve()

    # Under min_after_retrieve elements, so can not retrieve just yet
    _add_many(q, [1, 2])
    assert q.can_add()
    assert not q.can_retrieve()
    assert q.size == 2

    # Got to min_after_retrieve elements, can start retrieving
    _add_many(q, [3])
    assert q.can_retrieve()
    assert q.size == 3

    # But when we retrieve we are again under min_after_retrieve, so can not retrieve again
    _retrieve(q)
    assert not q.can_retrieve()
    assert q.size == 2

    # Getting back to the retrievable state with enough items in the buffer
    _add_many(q, [4, 5])
    assert q.can_add()
    assert q.can_retrieve()
    assert q.size == 4

    # Can overrun the capacity (as long as below extra_capacity), but can not add if we are above
    # shuffling_buffer_capacity
    _add_many(q, [6, 7, 8, 9])
    assert not q.can_add()
    with pytest.raises(RuntimeError):
        _add_many(q, [1])
    assert q.can_retrieve()
    assert q.size == 8

    # Getting one out. Still have more than shuffling_buffer_capacity
    _retrieve(q)
    assert not q.can_add()
    assert q.can_retrieve()
    assert q.size == 7

    # Retrieve enough to get back to addable state
    [_retrieve(q) for _ in range(4)]
    assert q.can_add()
    assert q.can_retrieve()
    assert q.size == 3

    # Retrieve the last element so we go under min_after_retrieve and can not retrieve any more
    _retrieve(q)
    assert q.can_add()
    assert not q.can_retrieve()
    with pytest.raises(RuntimeError):
        _retrieve(q)

    assert q.size == 2

    # finish() will allow us to deplete the buffer completely
    q.finish()
    assert not q.can_add()
    assert q.can_retrieve()
    assert q.size == 2

    _retrieve(q)
    assert not q.can_add()
    assert q.can_retrieve()
    assert q.size == 1

    _retrieve(q)
    assert not q.can_add()
    assert not q.can_retrieve()
    assert q.size == 0
Check can_add/can_retrieve flags at all possible states
petastorm/tests/test_shuffling_buffer.py