RichAsyncFunction extends AbstractRichFunction implements AsyncFunction<IN, OUT> { @Override public void setRuntimeContext(RuntimeContext runtimeContext) { Preconditions.checkNotNull(runtimeContext); if (runtimeContext instanceof IterationRuntimeContext) { super.setRuntimeContext( new RichAsyncFunctionIterationRuntimeContext( (IterationRuntimeContext) runtimeContext)); } else { super.setRuntimeContext(new RichAsyncFunctionRuntimeContext(runtimeContext)); } }  @Override void setRuntimeContext(RuntimeContext runtimeContext); @Override abstract void asyncInvoke(IN input, ResultFuture<OUT> resultFuture);  }
@Test public void testRuntimeContext() throws Exception { RichAsyncFunction<Integer, Integer> function = new RichAsyncFunction<Integer, Integer>() { private static final long serialVersionUID = 1707630162838967972L; @Override public void asyncInvoke(Integer input, ResultFuture<Integer> resultFuture) throws Exception { } }; final String taskName = "foobarTask"; final MetricGroup metricGroup = mock(MetricGroup.class); final int numberOfParallelSubtasks = 42; final int indexOfSubtask = 43; final int attemptNumber = 1337; final String taskNameWithSubtask = "barfoo"; final ExecutionConfig executionConfig = mock(ExecutionConfig.class); final ClassLoader userCodeClassLoader = mock(ClassLoader.class); RuntimeContext mockedRuntimeContext = mock(RuntimeContext.class); when(mockedRuntimeContext.getTaskName()).thenReturn(taskName); when(mockedRuntimeContext.getMetricGroup()).thenReturn(metricGroup); when(mockedRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(numberOfParallelSubtasks); when(mockedRuntimeContext.getIndexOfThisSubtask()).thenReturn(indexOfSubtask); when(mockedRuntimeContext.getAttemptNumber()).thenReturn(attemptNumber); when(mockedRuntimeContext.getTaskNameWithSubtasks()).thenReturn(taskNameWithSubtask); when(mockedRuntimeContext.getExecutionConfig()).thenReturn(executionConfig); when(mockedRuntimeContext.getUserCodeClassLoader()).thenReturn(userCodeClassLoader); function.setRuntimeContext(mockedRuntimeContext); RuntimeContext runtimeContext = function.getRuntimeContext(); assertEquals(taskName, runtimeContext.getTaskName()); assertEquals(metricGroup, runtimeContext.getMetricGroup()); assertEquals(numberOfParallelSubtasks, runtimeContext.getNumberOfParallelSubtasks()); assertEquals(indexOfSubtask, runtimeContext.getIndexOfThisSubtask()); assertEquals(attemptNumber, runtimeContext.getAttemptNumber()); assertEquals(taskNameWithSubtask, runtimeContext.getTaskNameWithSubtasks()); assertEquals(executionConfig, runtimeContext.getExecutionConfig()); assertEquals(userCodeClassLoader, runtimeContext.getUserCodeClassLoader()); try { runtimeContext.getDistributedCache(); fail("Expected getDistributedCached to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getState(new ValueStateDescriptor<>("foobar", Integer.class, 42)); fail("Expected getState to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getListState(new ListStateDescriptor<>("foobar", Integer.class)); fail("Expected getListState to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getReducingState(new ReducingStateDescriptor<>("foobar", new ReduceFunction<Integer>() { private static final long serialVersionUID = 2136425961884441050L; @Override public Integer reduce(Integer value1, Integer value2) throws Exception { return value1; } }, Integer.class)); fail("Expected getReducingState to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getAggregatingState(new AggregatingStateDescriptor<>("foobar", new AggregateFunction<Integer, Integer, Integer>() { @Override public Integer createAccumulator() { return null; } @Override public Integer add(Integer value, Integer accumulator) { return null; } @Override public Integer getResult(Integer accumulator) { return null; } @Override public Integer merge(Integer a, Integer b) { return null; } }, Integer.class)); } catch (UnsupportedOperationException e) { } try { runtimeContext.getFoldingState(new FoldingStateDescriptor<>("foobar", 0, new FoldFunction<Integer, Integer>() { @Override public Integer fold(Integer accumulator, Integer value) throws Exception { return accumulator; } }, Integer.class)); } catch (UnsupportedOperationException e) { } try { runtimeContext.getMapState(new MapStateDescriptor<>("foobar", Integer.class, String.class)); } catch (UnsupportedOperationException e) { } try { runtimeContext.addAccumulator("foobar", new Accumulator<Integer, Integer>() { private static final long serialVersionUID = -4673320336846482358L; @Override public void add(Integer value) { } @Override public Integer getLocalValue() { return null; } @Override public void resetLocal() { } @Override public void merge(Accumulator<Integer, Integer> other) { } @Override public Accumulator<Integer, Integer> clone() { return null; } }); fail("Expected addAccumulator to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getAccumulator("foobar"); fail("Expected getAccumulator to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getAllAccumulators(); fail("Expected getAllAccumulators to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getIntCounter("foobar"); fail("Expected getIntCounter to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getLongCounter("foobar"); fail("Expected getLongCounter to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getDoubleCounter("foobar"); fail("Expected getDoubleCounter to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getHistogram("foobar"); fail("Expected getHistogram to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.hasBroadcastVariable("foobar"); fail("Expected hasBroadcastVariable to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getBroadcastVariable("foobar"); fail("Expected getBroadcastVariable to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } try { runtimeContext.getBroadcastVariableWithInitializer("foobar", new BroadcastVariableInitializer<Object, Object>() { @Override public Object initializeBroadcastVariable(Iterable<Object> data) { return null; } }); fail("Expected getBroadcastVariableWithInitializer to fail with unsupported operation exception."); } catch (UnsupportedOperationException e) { } }