Examples with MockStreamingRuntimeContext org.apache.flink.streaming.util.MockStreamingRuntimeContext used on opensource projects

Example 1 with MockStreamingRuntimeContext

use of org.apache.flink.streaming.util.MockStreamingRuntimeContext in project flink by apache.

the class RMQSinkTest method testOverrideConnection.

@Test
public void testOverrideConnection() throws Exception {
    final Connection mockConnection = mock(Connection.class);
    Channel channel = mock(Channel.class);
    when(mockConnection.createChannel()).thenReturn(channel);
    RMQSink<String> rmqSink = new RMQSink<String>(rmqConnectionConfig, QUEUE_NAME, serializationSchema) {

        @Override
        protected Connection setupConnection() throws Exception {
            return mockConnection;
        }
    };
    rmqSink.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
    rmqSink.open(new Configuration());
    verify(mockConnection, times(1)).createChannel();
}

Example 2 with MockStreamingRuntimeContext

use of org.apache.flink.streaming.util.MockStreamingRuntimeContext in project flink by apache.

the class RMQSinkTest method createRMQSink.

private RMQSink<String> createRMQSink() throws Exception {
    RMQSink<String> rmqSink = new RMQSink<>(rmqConnectionConfig, QUEUE_NAME, serializationSchema);
    rmqSink.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
    rmqSink.open(new Configuration());
    return rmqSink;
}

Example 3 with MockStreamingRuntimeContext

use of org.apache.flink.streaming.util.MockStreamingRuntimeContext in project flink by apache.

the class WatermarkTrackerTest method test.

@Test
public void test() {
    long watermark = 0;
    TestWatermarkTracker ws = new TestWatermarkTracker();
    ws.open(new MockStreamingRuntimeContext(false, 1, 0));
    Assert.assertEquals(Long.MIN_VALUE, ws.updateWatermark(Long.MIN_VALUE));
    Assert.assertEquals(Long.MIN_VALUE, ws.updateWatermark(watermark));
    // timeout wm1
    clock.add(WatermarkTracker.DEFAULT_UPDATE_TIMEOUT_MILLIS + 1);
    Assert.assertEquals(watermark, ws.updateWatermark(watermark));
    Assert.assertEquals(watermark, ws.updateWatermark(watermark - 1));
    // min watermark
    wm1.watermark = watermark + 1;
    wm1.lastUpdated = clock.longValue();
    Assert.assertEquals(watermark, ws.updateWatermark(watermark));
    Assert.assertEquals(watermark + 1, ws.updateWatermark(watermark + 1));
}

Example 4 with MockStreamingRuntimeContext

use of org.apache.flink.streaming.util.MockStreamingRuntimeContext in project flink by apache.

the class JdbcDynamicTableSinkITCase method testFlushBufferWhenCheckpoint.

@Test
public void testFlushBufferWhenCheckpoint() throws Exception {
    Map<String, String> options = new HashMap<>();
    options.put("connector", "jdbc");
    options.put("url", DB_URL);
    options.put("table-name", OUTPUT_TABLE5);
    options.put("sink.buffer-flush.interval", "0");
    ResolvedSchema schema = ResolvedSchema.of(Column.physical("id", DataTypes.BIGINT().notNull()));
    DynamicTableSink tableSink = createTableSink(schema, options);
    SinkRuntimeProviderContext context = new SinkRuntimeProviderContext(false);
    SinkFunctionProvider sinkProvider = (SinkFunctionProvider) tableSink.getSinkRuntimeProvider(context);
    GenericJdbcSinkFunction<RowData> sinkFunction = (GenericJdbcSinkFunction<RowData>) sinkProvider.createSinkFunction();
    sinkFunction.setRuntimeContext(new MockStreamingRuntimeContext(true, 1, 0));
    sinkFunction.open(new Configuration());
    sinkFunction.invoke(GenericRowData.of(1L), SinkContextUtil.forTimestamp(1));
    sinkFunction.invoke(GenericRowData.of(2L), SinkContextUtil.forTimestamp(1));
    check(new Row[] {}, DB_URL, OUTPUT_TABLE5, new String[] { "id" });
    sinkFunction.snapshotState(new StateSnapshotContextSynchronousImpl(1, 1));
    check(new Row[] { Row.of(1L), Row.of(2L) }, DB_URL, OUTPUT_TABLE5, new String[] { "id" });
    sinkFunction.close();
}

Example 5 with MockStreamingRuntimeContext

use of org.apache.flink.streaming.util.MockStreamingRuntimeContext in project flink by apache.

the class BroadcastStateInputFormatTest method testReadBroadcastState.

@Test
public void testReadBroadcastState() throws Exception {
    try (TwoInputStreamOperatorTestHarness<Void, Integer, Void> testHarness = getTestHarness()) {
        testHarness.open();
        testHarness.processElement2(new StreamRecord<>(1));
        testHarness.processElement2(new StreamRecord<>(2));
        testHarness.processElement2(new StreamRecord<>(3));
        OperatorSubtaskState subtaskState = testHarness.snapshot(0, 0);
        OperatorState state = new OperatorState(OperatorIDGenerator.fromUid("uid"), 1, 4);
        state.putState(0, subtaskState);
        OperatorStateInputSplit split = new OperatorStateInputSplit(subtaskState.getManagedOperatorState(), 0);
        BroadcastStateInputFormat<Integer, Integer> format = new BroadcastStateInputFormat<>(state, new Configuration(), null, descriptor);
        format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
        format.open(split);
        Map<Integer, Integer> results = new HashMap<>(3);
        while (!format.reachedEnd()) {
            Tuple2<Integer, Integer> entry = format.nextRecord(null);
            results.put(entry.f0, entry.f1);
        }
        Map<Integer, Integer> expected = new HashMap<>(3);
        expected.put(1, 1);
        expected.put(2, 2);
        expected.put(3, 3);
        Assert.assertEquals("Failed to read correct list state from state backend", expected, results);
    }
}