def test_done_dataloader_iter_with_limit(mode, iterables, limit, num_fetches, expected):
    """Test that the `done` property for `dataloader_iter` gets set as expected."""
    loader = CombinedLoader(iterables, mode=mode)
    fetcher = _DataLoaderIterDataFetcher()
    loader.limits = limit
    fetcher.setup(loader)
    iter(fetcher)

    assert fetcher.done == (limit == 0)
    if num_fetches == 0:
        return

    dataloader_iter = next(fetcher)

    assert not dataloader_iter.done
    for _ in range(num_fetches):
        next(dataloader_iter)
    assert dataloader_iter.done == expected
    assert fetcher.done == expected

    if fetcher.done:
        with pytest.raises(StopIteration):
            next(dataloader_iter)
Test that the `done` property for `dataloader_iter` gets set as expected.
tests/tests_pytorch/loops/test_fetchers.py