Examples with IterationRuntimeContext org.apache.flink.api.common.functions.IterationRuntimeContext used on opensource projects

Example 1 with IterationRuntimeContext

use of org.apache.flink.api.common.functions.IterationRuntimeContext in project flink by apache.

the class RichAsyncFunctionTest method testIterationRuntimeContext.

/**
	 * Test the set of iteration runtime context methods in the context of a
	 * {@link RichAsyncFunction}.
	 */
@Test
public void testIterationRuntimeContext() throws Exception {
    RichAsyncFunction<Integer, Integer> function = new RichAsyncFunction<Integer, Integer>() {

        private static final long serialVersionUID = -2023923961609455894L;

        @Override
        public void asyncInvoke(Integer input, AsyncCollector<Integer> collector) throws Exception {
        // no op
        }
    };
    int superstepNumber = 42;
    IterationRuntimeContext mockedIterationRuntimeContext = mock(IterationRuntimeContext.class);
    when(mockedIterationRuntimeContext.getSuperstepNumber()).thenReturn(superstepNumber);
    function.setRuntimeContext(mockedIterationRuntimeContext);
    IterationRuntimeContext iterationRuntimeContext = function.getIterationRuntimeContext();
    assertEquals(superstepNumber, iterationRuntimeContext.getSuperstepNumber());
    try {
        iterationRuntimeContext.getIterationAggregator("foobar");
        fail("Expected getIterationAggregator to fail with unsupported operation exception");
    } catch (UnsupportedOperationException e) {
    // expected
    }
    try {
        iterationRuntimeContext.getPreviousIterationAggregate("foobar");
        fail("Expected getPreviousIterationAggregator to fail with unsupported operation exception");
    } catch (UnsupportedOperationException e) {
    // expected
    }
}