def test_get_chunks(chunks):
from unittest import mock
with mock.patch("dask_ml.utils.cpu_count", return_value=4):
result = check_chunks(n_samples=8000, n_features=4, chunks=chunks)
expected = (2000, 4)
assert result == expected
def test_get_chunks_raises():
with pytest.raises(AssertionError):
check_chunks(1, 1, chunks=(1, 2, 3))
with pytest.raises(AssertionError):
check_chunks(1, 1, chunks=[1, 2, 3])
with pytest.raises(ValueError):
check_chunks(1, 1, chunks=object())
def test_get_chunks_min(chunks):
result = check_chunks(n_samples=8, n_features=4, chunks=chunks)
expected = (100, 4)
assert result == expected
