@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));
        }
    }
@Test
    void testIterationRuntimeContext() {
        RichAsyncFunction<Integer, Integer> function =
                new RichAsyncFunction<Integer, Integer>() {
                    private static final long serialVersionUID = -2023923961609455894L;

                    @Override
                    public void asyncInvoke(Integer input, ResultFuture<Integer> resultFuture)
                            throws Exception {
                        // no op
                    }
                };

        int superstepNumber = 42;

        IterationRuntimeContext mockedIterationRuntimeContext = mock(IterationRuntimeContext.class);
        when(mockedIterationRuntimeContext.getSuperstepNumber()).thenReturn(superstepNumber);
        function.setRuntimeContext(mockedIterationRuntimeContext);

        IterationRuntimeContext iterationRuntimeContext = function.getIterationRuntimeContext();

        assertThat(iterationRuntimeContext.getSuperstepNumber()).isEqualTo(superstepNumber);

        assertThatThrownBy(() -> iterationRuntimeContext.getIterationAggregator("foobar"))
                .isInstanceOf(UnsupportedOperationException.class);

        assertThatThrownBy(() -> iterationRuntimeContext.getPreviousIterationAggregate("foobar"))
                .isInstanceOf(UnsupportedOperationException.class);
    }