@Override
    public void applyToAllPartitions(
            TwoOutputApplyPartitionFunction<OUT1, OUT2> applyPartitionFunction) throws Exception {
        if (isKeyed) {
            for (Object key : keySet) {
                partitionedContext
                        .getStateManager()
                        .executeInKeyContext(
                                () -> {
                                    try {
                                        applyPartitionFunction.apply(
                                                firstCollector,
                                                secondCollector,
                                                partitionedContext);
                                    } catch (Exception e) {
                                        throw new RuntimeException(e);
                                    }
                                },
                                key);
            }
        } else {
            // non-keyed operator has only one partition.
            applyPartitionFunction.apply(firstCollector, secondCollector, partitionedContext);
        }
    }
@Test
    void testApplyToAllPartitions() throws Exception {
        AtomicInteger counter = new AtomicInteger(0);
        List<Integer> collectedFromFirstOutput = new ArrayList<>();
        List<Long> collectedFromSecondOutput = new ArrayList<>();

        TestingTimestampCollector<Integer> firstCollector =
                TestingTimestampCollector.<Integer>builder()
                        .setCollectConsumer(collectedFromFirstOutput::add)
                        .build();
        TestingTimestampCollector<Long> secondCollector =
                TestingTimestampCollector.<Long>builder()
                        .setCollectConsumer(collectedFromSecondOutput::add)
                        .build();
        CompletableFuture<Void> cf = new CompletableFuture<>();
        StreamingRuntimeContext operatorRuntimeContext =
                ContextTestUtils.createStreamingRuntimeContext();
        DefaultRuntimeContext runtimeContext =
                new DefaultRuntimeContext(
                        operatorRuntimeContext.getJobInfo().getJobName(),
                        operatorRuntimeContext.getJobType(),
                        1,
                        2,
                        "mock-task",
                        operatorRuntimeContext.getMetricGroup());
        DefaultTwoOutputNonPartitionedContext<Integer, Long> nonPartitionedContext =
                new DefaultTwoOutputNonPartitionedContext<>(
                        runtimeContext,
                        new DefaultPartitionedContext(
                                runtimeContext,
                                Optional::empty,
                                (key) -> cf.complete(null),
                                UnsupportedProcessingTimeManager.INSTANCE,
                                ContextTestUtils.createStreamingRuntimeContext(),
                                new MockOperatorStateStore()),
                        firstCollector,
                        secondCollector,
                        false,
                        null);

        nonPartitionedContext.applyToAllPartitions(
                (firstOutput, secondOutput, ctx) -> {
                    counter.incrementAndGet();
                    firstOutput.collect(10);
                    secondOutput.collect(20L);
                });
        assertThat(counter.get()).isEqualTo(1);
        assertThat(cf).isNotCompleted();
        assertThat(collectedFromFirstOutput).containsExactly(10);
        assertThat(collectedFromSecondOutput).containsExactly(20L);
    }