def test_split(sample_collection, samples, batch_size, max_batch_size):
    """test splitting sample collection into batches"""

    fomi.SampleIdBatch.get_max_batch_size = mock.Mock(
        return_value=max_batch_size
    )

    #####
    result = fomi.SampleIdBatch.split(
        sample_collection, workers := 8, batch_size
    )
    #####

    expected_batch_size = min(
        (batch_size if batch_size is not None else len(samples) // workers),
        max_batch_size,
    )
    expected_batch_count = len(samples) // expected_batch_size

    start_and_stop_idxs = [
        (i * expected_batch_size, (i + 1) * expected_batch_size)
        for i in range(expected_batch_count)
    ]

    samples_ids = [sample.id for sample in samples]

    sample_collection.values.assert_called_once_with(
        "id", _enforce_natural_order=False
    )

    assert len(result) == expected_batch_count

    total = 0
    for idx, batch in enumerate(result):
        assert isinstance(batch, fomi.SampleIdBatch)
        assert batch.total == expected_batch_size
        total += batch.total

        start_idx, stop_idx = start_and_stop_idxs[idx]

        assert batch.sample_ids == tuple(samples_ids[start_idx:stop_idx])

    assert total == len(samples_ids)
test splitting sample collection into batches
tests/unittests/map/test_id_batch.py