def test_instance_slicing(
    monkeypatch,
    reader_class,
    world_size: Optional[int],
    num_workers: Optional[int],
    max_instances: Optional[int],
):
    """
    Ensure that the intances read by each worker are always unique and the total
    adds up to `max_instances`.
    """
    results: List[Set[int]] = []

    minimum_expected_result_size = max_instances or TOTAL_INSTANCES
    maximum_expected_result_size = max_instances or TOTAL_INSTANCES

    if world_size is not None and num_workers is not None:
        minimum_expected_result_size //= world_size
        minimum_expected_result_size //= num_workers
        maximum_expected_result_size = minimum_expected_result_size + 1
        for global_rank in range(world_size):
            monkeypatch.setattr(common_util, "is_distributed", lambda: True)
            monkeypatch.setattr(dist, "get_rank", lambda: global_rank)
            monkeypatch.setattr(dist, "get_world_size", lambda: world_size)
            for worker_id in range(num_workers):
                reader = reader_class(max_instances=max_instances)
                reader._set_worker_info(WorkerInfo(num_workers, worker_id))
                result = set(
                    x["index"].label for x in reader.read("the-path-doesnt-matter")  # type: ignore
                )
                results.append(result)
    elif world_size is not None:
        minimum_expected_result_size //= world_size
        maximum_expected_result_size = minimum_expected_result_size + 1
        for global_rank in range(world_size):
            monkeypatch.setattr(common_util, "is_distributed", lambda: True)
            monkeypatch.setattr(dist, "get_rank", lambda: global_rank)
            monkeypatch.setattr(dist, "get_world_size", lambda: world_size)
            reader = reader_class(max_instances=max_instances)
            result = set(
                x["index"].label for x in reader.read("the-path-doesnt-matter")  # type: ignore
            )
            results.append(result)
    elif num_workers is not None:
        minimum_expected_result_size //= num_workers
        maximum_expected_result_size = minimum_expected_result_size + 1
        for worker_id in range(num_workers):
            reader = reader_class(max_instances=max_instances)
            reader._set_worker_info(WorkerInfo(num_workers, worker_id))
            result = set(
                x["index"].label for x in reader.read("the-path-doesnt-matter")  # type: ignore
            )
            results.append(result)
    else:
        reader = reader_class(max_instances=max_instances)
        result = set(
            x["index"].label for x in reader.read("the-path-doesnt-matter")  # type: ignore
        )
        results.append(result)

    # We need to check that all of the result sets are mutually exclusive and that they're
    # union has size `max_instances`.
    # Checking that they're mutually exclusive is equivalent to checking that the sum
    # of the size of each set is equal to the size of the union.

    union: Set[int] = set()
    total: int = 0
    for result in results:
        union |= result
        total += len(result)
        # Also make sure the size of the set is within the expected bounds.
        assert minimum_expected_result_size <= len(result)
        assert len(result) <= maximum_expected_result_size

    assert len(union) == total == (max_instances or TOTAL_INSTANCES)
Ensure that the intances read by each worker are always unique and the total
adds up to `max_instances`.
tests/data/dataset_readers/dataset_reader_test.py