Can the test count be set when using a TestNG DataProvider that returns an Iterator

Question

Is there a way to set the number of records expected, while still returning an iterator instead of an object array?

I trying to use TestNG for regression testing. I have a data source with many test cases. I would like to use an iterator to provide these to the test methods, but without the array the number of expected tests is not displayed in IntelliJ.

I have tried setting different properties in the ITestContext, ITestNGMethod, and Method, but nothing seems to work. For example, I am trying to do something like:

@DataProvider(name = "TestCaseProvider")
public static Iterator<Object[]> TestCases (ITestContext context)
{
    MyTestCaseHeader header = new MyTestCaseHeader(...);

    context.setExpectedInvocationCount(header.getTestCaseCount());

    return header.getIterator();
}