def test_buffer(capacity, chunk_len, sample_shape):
"""Builds a Buffer with the provided `capacity` and insert `capacity * 3`
samples into the buffer in chunks of shape `(chunk_len,) + sample_shape`.
We always insert chunks with consecutive integers.
* `len(buffer)` should increase until we reach capacity.
* `buffer._idx` should loop between 0 and `capacity - 1`.
* After every insertion, samples should be in expected range, verifying
FIFO insertion.
* Mutating the inserted chunk shouldn't mutate the buffer.
"""
buf = Buffer(capacity,
sample_shapes={'a': sample_shape, 'b': sample_shape},
dtypes={'a': float, 'b': float})
to_insert = 3 * capacity
for i in range(0, to_insert, chunk_len):
assert len(buf) == min(i, capacity)
assert buf._idx == i % capacity
chunk_a = _fill_chunk(i, chunk_len, sample_shape)
chunk_b = _fill_chunk(i + to_insert, chunk_len, sample_shape)
buf.store({'a': chunk_a, 'b': chunk_b})
samples = buf.sample(100)
assert set(samples.keys()) == {'a', 'b'}, samples.keys()
_check_bound(i + chunk_len, capacity, samples['a'])
_check_bound(i + chunk_len + to_insert, capacity, samples['b'])
assert np.all(samples['b'] - samples['a'] == to_insert)
# Confirm that buffer is not mutable from inserted sample.
chunk_a[:] = np.nan
chunk_b[:] = np.nan
assert not np.any(np.isnan(buf._arrays['a']))
assert not np.any(np.isnan(buf._arrays['b']))
def test_buffer_from_data():
data = np.ndarray([50, 30], dtype=bool)
buf = Buffer.from_data({'k': data})
assert buf._arrays['k'] is not data
assert data.dtype == buf._arrays['k'].dtype
assert np.array_equal(buf._arrays['k'], data)
def test_buffer_init_errors():
with pytest.raises(KeyError, match=r"sample_shape and dtypes.*"):
Buffer(10, dict(a=(2, 1), b=(3, )), dtypes=dict(a="float32", c=bool))
def test_buffer_sample_errors():
b = Buffer(10, {'k': (2, 1)}, dtypes={'k': bool})
with pytest.raises(ValueError):
b.sample(5)
def buf():
return Buffer(capacity, {'k': sample_shape}, {'k': dtype})