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