[ML] Test DatasetAffinityFunctionWrapperTest failed with UnnecessaryStubbingException

NOTE: This is not reproduced locally, but reproduced on TC

org.mockito.exceptions.misusing.UnnecessaryStubbingException: Unnecessary stubbings detected in test class: DatasetAffinityFunctionWrapperTest Clean & maintainable test code requires zero unnecessary code. Following stubbings are unnecessary (click to navigate to relevant line of code): 1. -> at org.apache.ignite.ml.dataset.impl.cache.util.DatasetAffinityFunctionWrapperTest.testPartition(DatasetAffinityFunctionWrapperTest.java:80) Please remove unnecessary stubbings or use 'lenient' strictness. More info: javadoc for UnnecessaryStubbingException class.
org.mockito.exceptions.misusing.UnnecessaryStubbingException:
Unnecessary stubbings detected in test class: DatasetAffinityFunctionWrapperTest
Clean & maintainable test code requires zero unnecessary code.
Following stubbings are unnecessary (click to navigate to relevant line of code):
1. -> at org.apache.ignite.ml.dataset.impl.cache.util.DatasetAffinityFunctionWrapperTest.testPartition(DatasetAffinityFunctionWrapperTest.java:80)
Please remove unnecessary stubbings or use 'lenient' strictness. More info: javadoc for UnnecessaryStubbingException class.