public static <T extends Serializable> List<T> runSourceFunction(SourceFunction<T> sourceFunction) throws Exception { if (sourceFunction instanceof RichFunction) { return runRichSourceFunction(sourceFunction); } else { return runNonRichSourceFunction(sourceFunction); } }
@Test public void fromCollectionTest() throws Exception { List<Integer> expectedList = Arrays.asList(1, 2, 3); List<Integer> actualList = SourceFunctionUtil.runSourceFunction( CommonTestUtils.createCopySerializable(new FromElementsFunction<Integer>( IntSerializer.INSTANCE, Arrays.asList(1, 2, 3)))); assertEquals(expectedList, actualList); } }
private static <T extends Serializable> List<T> runRichSourceFunction(SourceFunction<T> sourceFunction) throws Exception { try (MockEnvironment environment = new MockEnvironmentBuilder() .setTaskName("MockTask") .setMemorySize(3 * 1024 * 1024) .setInputSplitProvider(new MockInputSplitProvider()) .setBufferSize(1024) .build()) { AbstractStreamOperator<?> operator = mock(AbstractStreamOperator.class); when(operator.getExecutionConfig()).thenReturn(new ExecutionConfig()); RuntimeContext runtimeContext = new StreamingRuntimeContext( operator, environment, new HashMap<>()); ((RichFunction) sourceFunction).setRuntimeContext(runtimeContext); ((RichFunction) sourceFunction).open(new Configuration()); return runNonRichSourceFunction(sourceFunction); } }
@Test public void fromElementsTest() throws Exception { List<Integer> expectedList = Arrays.asList(1, 2, 3); List<Integer> actualList = SourceFunctionUtil.runSourceFunction(CommonTestUtils.createCopySerializable( new FromElementsFunction<Integer>( IntSerializer.INSTANCE, 1, 2, 3))); assertEquals(expectedList, actualList); }