datasets

1
import pytest
2

3
import datasets.config
4
from datasets.utils.info_utils import is_small_dataset
5

6

7
@pytest.mark.parametrize("dataset_size", [None, 400 * 2**20, 600 * 2**20])
8
@pytest.mark.parametrize("input_in_memory_max_size", ["default", 0, 100 * 2**20, 900 * 2**20])
9
def test_is_small_dataset(dataset_size, input_in_memory_max_size, monkeypatch):
10
    if input_in_memory_max_size != "default":
11
        monkeypatch.setattr(datasets.config, "IN_MEMORY_MAX_SIZE", input_in_memory_max_size)
12
    in_memory_max_size = datasets.config.IN_MEMORY_MAX_SIZE
13
    if input_in_memory_max_size == "default":
14
        assert in_memory_max_size == 0
15
    else:
16
        assert in_memory_max_size == input_in_memory_max_size
17
    if dataset_size and in_memory_max_size:
18
        expected = dataset_size < in_memory_max_size
19
    else:
20
        expected = False
21
    result = is_small_dataset(dataset_size)
22
    assert result == expected
23