Search in sources :

Example 6 with RuntimeContext

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

the class InternalWindowFunctionTest method testInternalIterableProcessAllWindowFunction.

@SuppressWarnings("unchecked")
@Test
public void testInternalIterableProcessAllWindowFunction() throws Exception {
    ProcessAllWindowFunctionMock mock = mock(ProcessAllWindowFunctionMock.class);
    InternalIterableProcessAllWindowFunction<Long, String, TimeWindow> windowFunction = new InternalIterableProcessAllWindowFunction<>(mock);
    // check setOutputType
    TypeInformation<String> stringType = BasicTypeInfo.STRING_TYPE_INFO;
    ExecutionConfig execConf = new ExecutionConfig();
    execConf.setParallelism(42);
    StreamingFunctionUtils.setOutputType(windowFunction, stringType, execConf);
    verify(mock).setOutputType(stringType, execConf);
    // check open
    Configuration config = new Configuration();
    windowFunction.open(config);
    verify(mock).open(config);
    // check setRuntimeContext
    RuntimeContext rCtx = mock(RuntimeContext.class);
    windowFunction.setRuntimeContext(rCtx);
    verify(mock).setRuntimeContext(rCtx);
    // check apply
    TimeWindow w = mock(TimeWindow.class);
    Iterable<Long> i = (Iterable<Long>) mock(Iterable.class);
    Collector<String> c = (Collector<String>) mock(Collector.class);
    windowFunction.apply(((byte) 0), w, i, c);
    verify(mock).process((ProcessAllWindowFunctionMock.Context) anyObject(), eq(i), eq(c));
    // check close
    windowFunction.close();
    verify(mock).close();
}
Also used : Configuration(org.apache.flink.configuration.Configuration) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) TimeWindow(org.apache.flink.streaming.api.windowing.windows.TimeWindow) Collector(org.apache.flink.util.Collector) InternalIterableProcessAllWindowFunction(org.apache.flink.streaming.runtime.operators.windowing.functions.InternalIterableProcessAllWindowFunction) RuntimeContext(org.apache.flink.api.common.functions.RuntimeContext) Test(org.junit.Test)

Example 7 with RuntimeContext

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

the class InternalWindowFunctionTest method testInternalSingleValueAllWindowFunction.

@SuppressWarnings("unchecked")
@Test
public void testInternalSingleValueAllWindowFunction() throws Exception {
    AllWindowFunctionMock mock = mock(AllWindowFunctionMock.class);
    InternalSingleValueAllWindowFunction<Long, String, TimeWindow> windowFunction = new InternalSingleValueAllWindowFunction<>(mock);
    // check setOutputType
    TypeInformation<String> stringType = BasicTypeInfo.STRING_TYPE_INFO;
    ExecutionConfig execConf = new ExecutionConfig();
    execConf.setParallelism(42);
    StreamingFunctionUtils.setOutputType(windowFunction, stringType, execConf);
    verify(mock).setOutputType(stringType, execConf);
    // check open
    Configuration config = new Configuration();
    windowFunction.open(config);
    verify(mock).open(config);
    // check setRuntimeContext
    RuntimeContext rCtx = mock(RuntimeContext.class);
    windowFunction.setRuntimeContext(rCtx);
    verify(mock).setRuntimeContext(rCtx);
    // check apply
    TimeWindow w = mock(TimeWindow.class);
    Collector<String> c = (Collector<String>) mock(Collector.class);
    windowFunction.apply(((byte) 0), w, 23L, c);
    verify(mock).apply(eq(w), (Iterable<Long>) argThat(IsIterableContainingInOrder.contains(23L)), eq(c));
    // check close
    windowFunction.close();
    verify(mock).close();
}
Also used : InternalSingleValueAllWindowFunction(org.apache.flink.streaming.runtime.operators.windowing.functions.InternalSingleValueAllWindowFunction) Configuration(org.apache.flink.configuration.Configuration) Collector(org.apache.flink.util.Collector) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) RuntimeContext(org.apache.flink.api.common.functions.RuntimeContext) TimeWindow(org.apache.flink.streaming.api.windowing.windows.TimeWindow) Test(org.junit.Test)

Example 8 with RuntimeContext

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

the class InternalWindowFunctionTest method testInternalSingleValueProcessAllWindowFunction.

@SuppressWarnings("unchecked")
@Test
public void testInternalSingleValueProcessAllWindowFunction() throws Exception {
    ProcessAllWindowFunctionMock mock = mock(ProcessAllWindowFunctionMock.class);
    InternalSingleValueProcessAllWindowFunction<Long, String, TimeWindow> windowFunction = new InternalSingleValueProcessAllWindowFunction<>(mock);
    // check setOutputType
    TypeInformation<String> stringType = BasicTypeInfo.STRING_TYPE_INFO;
    ExecutionConfig execConf = new ExecutionConfig();
    execConf.setParallelism(42);
    StreamingFunctionUtils.setOutputType(windowFunction, stringType, execConf);
    verify(mock).setOutputType(stringType, execConf);
    // check open
    Configuration config = new Configuration();
    windowFunction.open(config);
    verify(mock).open(config);
    // check setRuntimeContext
    RuntimeContext rCtx = mock(RuntimeContext.class);
    windowFunction.setRuntimeContext(rCtx);
    verify(mock).setRuntimeContext(rCtx);
    // check apply
    TimeWindow w = mock(TimeWindow.class);
    Collector<String> c = (Collector<String>) mock(Collector.class);
    windowFunction.apply(((byte) 0), w, 23L, c);
    verify(mock).process((ProcessAllWindowFunctionMock.Context) anyObject(), (Iterable<Long>) argThat(IsIterableContainingInOrder.contains(23L)), eq(c));
    // check close
    windowFunction.close();
    verify(mock).close();
}
Also used : Configuration(org.apache.flink.configuration.Configuration) Collector(org.apache.flink.util.Collector) InternalSingleValueProcessAllWindowFunction(org.apache.flink.streaming.runtime.operators.windowing.functions.InternalSingleValueProcessAllWindowFunction) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) RuntimeContext(org.apache.flink.api.common.functions.RuntimeContext) TimeWindow(org.apache.flink.streaming.api.windowing.windows.TimeWindow) Test(org.junit.Test)

Example 9 with RuntimeContext

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

the class SourceFunctionUtil method runSourceFunction.

public static <T extends Serializable> List<T> runSourceFunction(SourceFunction<T> sourceFunction) throws Exception {
    final List<T> outputs = new ArrayList<T>();
    if (sourceFunction instanceof RichFunction) {
        AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class);
        when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig());
        RuntimeContext runtimeContext = new StreamingRuntimeContext(operator, new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024), new HashMap<String, Accumulator<?, ?>>());
        ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext);
        ((RichFunction) sourceFunction).open(new Configuration());
    }
    try {
        SourceFunction.SourceContext<T> ctx = new CollectingSourceContext<T>(new Object(), outputs);
        sourceFunction.run(ctx);
    } catch (Exception e) {
        throw new RuntimeException("Cannot invoke source.", e);
    }
    return outputs;
}
Also used : Accumulator(org.apache.flink.api.common.accumulators.Accumulator) SourceFunction(org.apache.flink.streaming.api.functions.source.SourceFunction) StreamingRuntimeContext(org.apache.flink.streaming.api.operators.StreamingRuntimeContext) Configuration(org.apache.flink.configuration.Configuration) RichFunction(org.apache.flink.api.common.functions.RichFunction) ArrayList(java.util.ArrayList) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) MockEnvironment(org.apache.flink.runtime.operators.testutils.MockEnvironment) RuntimeContext(org.apache.flink.api.common.functions.RuntimeContext) StreamingRuntimeContext(org.apache.flink.streaming.api.operators.StreamingRuntimeContext) MockInputSplitProvider(org.apache.flink.runtime.operators.testutils.MockInputSplitProvider)

Example 10 with RuntimeContext

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

the class ReduceOperatorTest method testReduceCollectionWithRuntimeContext.

@Test
public void testReduceCollectionWithRuntimeContext() {
    try {
        final String taskName = "Test Task";
        final AtomicBoolean opened = new AtomicBoolean();
        final AtomicBoolean closed = new AtomicBoolean();
        final ReduceFunction<Tuple2<String, Integer>> reducer = new RichReduceFunction<Tuple2<String, Integer>>() {

            @Override
            public Tuple2<String, Integer> reduce(Tuple2<String, Integer> value1, Tuple2<String, Integer> value2) throws Exception {
                return new Tuple2<String, Integer>(value1.f0, value1.f1 + value2.f1);
            }

            @Override
            public void open(Configuration parameters) throws Exception {
                opened.set(true);
                RuntimeContext ctx = getRuntimeContext();
                assertEquals(0, ctx.getIndexOfThisSubtask());
                assertEquals(1, ctx.getNumberOfParallelSubtasks());
                assertEquals(taskName, ctx.getTaskName());
            }

            @Override
            public void close() throws Exception {
                closed.set(true);
            }
        };
        ReduceOperatorBase<Tuple2<String, Integer>, ReduceFunction<Tuple2<String, Integer>>> op = new ReduceOperatorBase<Tuple2<String, Integer>, ReduceFunction<Tuple2<String, Integer>>>(reducer, new UnaryOperatorInformation<Tuple2<String, Integer>, Tuple2<String, Integer>>(TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, Integer>"), TypeInfoParser.<Tuple2<String, Integer>>parse("Tuple2<String, " + "Integer>")), new int[] { 0 }, "TestReducer");
        List<Tuple2<String, Integer>> input = new ArrayList<Tuple2<String, Integer>>(asList(new Tuple2<String, Integer>("foo", 1), new Tuple2<String, Integer>("foo", 3), new Tuple2<String, Integer>("bar", 2), new Tuple2<String, Integer>("bar", 4)));
        final TaskInfo taskInfo = new TaskInfo(taskName, 1, 0, 1, 0);
        ExecutionConfig executionConfig = new ExecutionConfig();
        executionConfig.disableObjectReuse();
        List<Tuple2<String, Integer>> resultMutableSafe = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        executionConfig.enableObjectReuse();
        List<Tuple2<String, Integer>> resultRegular = op.executeOnCollections(input, new RuntimeUDFContext(taskInfo, null, executionConfig, new HashMap<String, Future<Path>>(), new HashMap<String, Accumulator<?, ?>>(), new UnregisteredMetricsGroup()), executionConfig);
        Set<Tuple2<String, Integer>> resultSetMutableSafe = new HashSet<Tuple2<String, Integer>>(resultMutableSafe);
        Set<Tuple2<String, Integer>> resultSetRegular = new HashSet<Tuple2<String, Integer>>(resultRegular);
        Set<Tuple2<String, Integer>> expectedResult = new HashSet<Tuple2<String, Integer>>(asList(new Tuple2<String, Integer>("foo", 4), new Tuple2<String, Integer>("bar", 6)));
        assertEquals(expectedResult, resultSetMutableSafe);
        assertEquals(expectedResult, resultSetRegular);
        assertTrue(opened.get());
        assertTrue(closed.get());
    } catch (Exception e) {
        e.printStackTrace();
        fail(e.getMessage());
    }
}
Also used : Path(org.apache.flink.core.fs.Path) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) Configuration(org.apache.flink.configuration.Configuration) HashMap(java.util.HashMap) ReduceFunction(org.apache.flink.api.common.functions.ReduceFunction) RichReduceFunction(org.apache.flink.api.common.functions.RichReduceFunction) ArrayList(java.util.ArrayList) ExecutionConfig(org.apache.flink.api.common.ExecutionConfig) TaskInfo(org.apache.flink.api.common.TaskInfo) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) RichReduceFunction(org.apache.flink.api.common.functions.RichReduceFunction) Tuple2(org.apache.flink.api.java.tuple.Tuple2) RuntimeUDFContext(org.apache.flink.api.common.functions.util.RuntimeUDFContext) RuntimeContext(org.apache.flink.api.common.functions.RuntimeContext) HashSet(java.util.HashSet) Test(org.junit.Test)

Aggregations

RuntimeContext (org.apache.flink.api.common.functions.RuntimeContext)30 ExecutionConfig (org.apache.flink.api.common.ExecutionConfig)19 Test (org.junit.Test)17 Configuration (org.apache.flink.configuration.Configuration)16 Collector (org.apache.flink.util.Collector)14 ArrayList (java.util.ArrayList)10 HashMap (java.util.HashMap)10 TimeWindow (org.apache.flink.streaming.api.windowing.windows.TimeWindow)10 HashSet (java.util.HashSet)5 TaskInfo (org.apache.flink.api.common.TaskInfo)5 RuntimeUDFContext (org.apache.flink.api.common.functions.util.RuntimeUDFContext)5 UnregisteredMetricsGroup (org.apache.flink.metrics.groups.UnregisteredMetricsGroup)5 AtomicBoolean (java.util.concurrent.atomic.AtomicBoolean)4 Accumulator (org.apache.flink.api.common.accumulators.Accumulator)4 StreamExecutionEnvironment (org.apache.flink.streaming.api.environment.StreamExecutionEnvironment)4 Map (java.util.Map)3 Set (java.util.Set)3 Path (org.apache.flink.core.fs.Path)3 StreamingRuntimeContext (org.apache.flink.streaming.api.operators.StreamingRuntimeContext)3 InetSocketAddress (java.net.InetSocketAddress)2