Flink test for broadcast state

Test.kt

@Test
fun `snapshot and subsequent restore`() {
    // non-empty function (knows about active customers)
    val function = getFunctionForTesting()
    val (operatorId, stateFinalizer) = getHarness(function).use { harness ->
        harness.processBroadcastElement(GenericChangeMessage(), 1L)
        val stateFinalizer = harness.snapshotWithLocalState(0L, 1L)
        assertEquals(1, stateFinalizer.jobManagerOwnedState.managedOperatorState.size)
        Pair(harness.streamConfig.operatorID, stateFinalizer)
    }

    // call constructor explicitly because ProcessFunctionTestHarnesses.forBroadcastProcessFunction calls
    // open with empty state, and we don't want that
    BroadcastOperatorTestHarness(
        CoBroadcastWithNonKeyedOperator(
            // empty function to make sure it reads from state
            MyFunction(emptyList()),
            listOf(ServicesCacheFlinkBridge.BROADCAST_STATE_DESCRIPTOR)
        ),
        1,
        1,
        0
    ).use { harness ->
        harness.streamConfig.operatorID = operatorId

        // unfortunately the API doesn't offer a way to set attempt number, so for now we use reflection
        harness.environment.javaClass.getDeclaredField("taskInfo").apply {
            isAccessible = true
            val taskInfo = with(harness.environment.taskInfo) {
                TaskInfo(
                    taskName,
                    maxNumberOfParallelSubtasks,
                    indexOfThisSubtask,
                    numberOfParallelSubtasks,
                    1,
                    allocationIDAsString
                )
            }
            set(harness.environment, taskInfo)
        }

        harness.initializeState(stateFinalizer.jobManagerOwnedState)
        harness.open()

        val event = getEvent()
        harness.processElement(event, 1L)
        assertEquals(4, harness.output.size)
    }
}