背景
我们经常需要对KeyedBroadcastProcessFunction函数进行单元测试,以确保上线之前这个函数的功能是正常的,包括里面的广播状态和键值分区状态
测试KeyedBroadcastProcessFunction类
java
@Test
public void testHarnessForKeyedBroadcastProcessFunction() throws Exception {
KeyedBroadcastProcessFunction<String, String, String, String> function = new MyKeyedBroadcastProcessFunction();
// 键值分区状态
final ValueStateDescriptor<String> valueStateDescriptor =
new ValueStateDescriptor<>("item", BasicTypeInfo.STRING_TYPE_INFO);
// 广播状态
final MapStateDescriptor<String, String> ruleStateDescriptor = new MapStateDescriptor<>("RulesBroadcastState",
BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
KeyedBroadcastOperatorTestHarness<String, String, String, String> harness =
ProcessFunctionTestHarnesses.forKeyedBroadcastProcessFunction(function, x -> x,
TypeInformation.of(String.class), ruleStateDescriptor);
harness.processBroadcastElement("0", 1);
harness.processBroadcastElement("000", 2);
harness.processElement("1", 10);
// 判断键值分区状态(注意这里最好就只是某个key下面,也就是分组key直接设置为x->"固定常数值"即可)
ValueState<String> valueState = function.getRuntimeContext().getState(valueStateDescriptor);
Assert.assertEquals(valueState.value(), "1");
// 判断广播状态
BroadcastState<String, String> broadcastState = harness.getBroadcastState(ruleStateDescriptor);
Assert.assertTrue(broadcastState.contains("0"));
Assert.assertTrue(broadcastState.contains("000"));
// 判断输出的列表
Assert.assertEquals(harness.extractOutputValues(), Arrays.asList("0", "000", "1"));
}
关键代码:
1.获取键值分区状态
java
ValueState<String> valueState = function.getRuntimeContext().getState(valueStateDescriptor);
2.获取广播状态:
java
BroadcastState<String, String> broadcastState = harness.getBroadcastState(ruleStateDescriptor);
3.工具类
java
public class ProcessFunctionTestHarnesses {
public ProcessFunctionTestHarnesses() {
}
public static <IN, OUT> OneInputStreamOperatorTestHarness<IN, OUT> forProcessFunction(ProcessFunction<IN, OUT> function) throws Exception {
OneInputStreamOperatorTestHarness<IN, OUT> testHarness = new OneInputStreamOperatorTestHarness(new ProcessOperator((ProcessFunction)Preconditions.checkNotNull(function)), 1, 1, 0);
testHarness.setup();
testHarness.open();
return testHarness;
}
public static <K, IN, OUT> KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> forKeyedProcessFunction(KeyedProcessFunction<K, IN, OUT> function, KeySelector<IN, K> keySelector, TypeInformation<K> keyType) throws Exception {
KeyedOneInputStreamOperatorTestHarness<K, IN, OUT> testHarness = new KeyedOneInputStreamOperatorTestHarness(new KeyedProcessOperator((KeyedProcessFunction)Preconditions.checkNotNull(function)), keySelector, keyType, 1, 1, 0);
testHarness.open();
return testHarness;
}
public static <IN1, IN2, OUT> TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> forCoProcessFunction(CoProcessFunction<IN1, IN2, OUT> function) throws Exception {
TwoInputStreamOperatorTestHarness<IN1, IN2, OUT> testHarness = new TwoInputStreamOperatorTestHarness(new CoProcessOperator((CoProcessFunction)Preconditions.checkNotNull(function)), 1, 1, 0);
testHarness.open();
return testHarness;
}
public static <K, IN1, IN2, OUT> KeyedTwoInputStreamOperatorTestHarness<K, IN1, IN2, OUT> forKeyedCoProcessFunction(KeyedCoProcessFunction<K, IN1, IN2, OUT> function, KeySelector<IN1, K> keySelector1, KeySelector<IN2, K> keySelector2, TypeInformation<K> keyType) throws Exception {
KeyedTwoInputStreamOperatorTestHarness<K, IN1, IN2, OUT> testHarness = new KeyedTwoInputStreamOperatorTestHarness(new KeyedCoProcessOperator((KeyedCoProcessFunction)Preconditions.checkNotNull(function)), keySelector1, keySelector2, keyType, 1, 1, 0);
testHarness.open();
return testHarness;
}
public static <IN1, IN2, OUT> BroadcastOperatorTestHarness<IN1, IN2, OUT> forBroadcastProcessFunction(BroadcastProcessFunction<IN1, IN2, OUT> function, MapStateDescriptor<?, ?>... descriptors) throws Exception {
BroadcastOperatorTestHarness<IN1, IN2, OUT> testHarness = new BroadcastOperatorTestHarness(new CoBroadcastWithNonKeyedOperator((BroadcastProcessFunction)Preconditions.checkNotNull(function), Arrays.asList(descriptors)), 1, 1, 0);
testHarness.open();
return testHarness;
}
public static <K, IN1, IN2, OUT> KeyedBroadcastOperatorTestHarness<K, IN1, IN2, OUT> forKeyedBroadcastProcessFunction(KeyedBroadcastProcessFunction<K, IN1, IN2, OUT> function, KeySelector<IN1, K> keySelector, TypeInformation<K> keyType, MapStateDescriptor<?, ?>... descriptors) throws Exception {
KeyedBroadcastOperatorTestHarness<K, IN1, IN2, OUT> testHarness = new KeyedBroadcastOperatorTestHarness(new CoBroadcastWithKeyedOperator((KeyedBroadcastProcessFunction)Preconditions.checkNotNull(function), Arrays.asList(descriptors)), keySelector, keyType, 1, 1, 0);
testHarness.open();
return testHarness;
}
}