use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit 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);
}
}
use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit in project flink by apache.
the class ListStateInputFormatTest method testReadListOperatorState.
@Test
public void testReadListOperatorState() throws Exception {
try (OneInputStreamOperatorTestHarness<Integer, Void> testHarness = getTestHarness()) {
testHarness.open();
testHarness.processElement(1, 0);
testHarness.processElement(2, 0);
testHarness.processElement(3, 0);
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);
ListStateInputFormat<Integer> format = new ListStateInputFormat<>(state, new Configuration(), null, descriptor);
format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
format.open(split);
List<Integer> results = new ArrayList<>();
while (!format.reachedEnd()) {
results.add(format.nextRecord(0));
}
results.sort(Comparator.naturalOrder());
Assert.assertEquals("Failed to read correct list state from state backend", Arrays.asList(1, 2, 3), results);
}
}
use of org.apache.flink.state.api.input.splits.OperatorStateInputSplit in project flink by apache.
the class UnionStateInputFormatTest method testReadUnionOperatorState.
@Test
public void testReadUnionOperatorState() throws Exception {
try (OneInputStreamOperatorTestHarness<Integer, Void> testHarness = getTestHarness()) {
testHarness.open();
testHarness.processElement(1, 0);
testHarness.processElement(2, 0);
testHarness.processElement(3, 0);
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);
UnionStateInputFormat<Integer> format = new UnionStateInputFormat<>(state, new Configuration(), null, descriptor);
format.setRuntimeContext(new MockStreamingRuntimeContext(false, 1, 0));
format.open(split);
List<Integer> results = new ArrayList<>();
while (!format.reachedEnd()) {
results.add(format.nextRecord(0));
}
results.sort(Comparator.naturalOrder());
Assert.assertEquals("Failed to read correct list state from state backend", Arrays.asList(1, 2, 3), results);
}
}
Aggregations