Repository: apache/flink
File: flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/aggregate/MiniBatchGroupAggFunctionTest.java (Java)

Code before:
        // Verify the output contains results for both "bbb" and "ccc"
        assertThat(output.stream().anyMatch(row -> row.getString(0).toString().equals("bbb")))
                .isTrue();
        assertThat(output.stream().anyMatch(row -> row.getString(0).toString().equals("ccc")))
                .isTrue();
    }

    /** Test collector that stores output rows. */
    private static class TestCollector implements Collector<RowData> {
        private final List<RowData> output;

        TestCollector(List<RowData> output) {
            this.output = output;
        }

        @Override
        public void collect(RowData record) {
            // Make a copy since the record might be reused
            GenericRowData copy = GenericRowData.of(record.getString(0), record.getLong(1));
            copy.setRowKind(record.getRowKind());
            output.add(copy);
        }

        @Override
        public void close() {}
    }

    /**
     * A test version of MiniBatchGroupAggFunction that doesn't require state infrastructure. It
     * simulates the behavior of the real function but uses in-memory maps instead of Flink state.
     */
    private static class TestMiniBatchGroupAggFunction {
        private final Map<RowData, RowData> accumulators = new LinkedHashMap<>();
        private RowData currentKey;

        public void finishBundle(Map<RowData, List<RowData>> buffer, Collector<RowData> out)
                throws Exception {
            for (Map.Entry<RowData, List<RowData>> entry : buffer.entrySet()) {
                RowData currentKey = entry.getKey();
                List<RowData> inputRows = entry.getValue();

                boolean firstRow = false;

                // Simulate: set current key and get accumulator from state
                this.currentKey = currentKey;
                RowData acc = accumulators.get(currentKey);

                if (acc == null) {
                    // Don't create a new accumulator for a retraction message.
                    java.util.Iterator<RowData> inputIter = inputRows.iterator();
                    while (inputIter.hasNext()) {

Review comment (suggestion):
You don't need to reimplement test output collectors. Almost for sure you can use one of the existing ones.

Code after:

    private MiniBatchGroupAggFunction createFunction(boolean generateUpdateBefore) {
        return new MiniBatchGroupAggFunction(
                genAggsHandler,
                genRecordEqualiser,
                accTypes,
                inputRowType,
                -1, // no COUNT(*) for this test
                generateUpdateBefore,
                0); // no state retention
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private OneInputStreamOperatorTestHarness<RowData, RowData> createTestHarness(
            MiniBatchGroupAggFunction function) throws Exception {
        CountBundleTrigger trigger = new CountBundleTrigger<>(10);
        KeyedMapBundleOperator operator = new KeyedMapBundleOperator(function, trigger);
        return new KeyedOneInputStreamOperatorTestHarness<>(
                operator, keySelector, keySelector.getProducedType());
    }

    /**
     * Verifies that when finishBundle processes a key with only retraction messages (which gets
     * filtered out because there's no accumulator state), the method continues to process
     * subsequent keys in the bundle instead of returning early.
     */
    @Test
    void testFinishBundleContinuesAfterEmptyInputRows() throws Exception {
        MiniBatchGroupAggFunction function = createFunction(false);
        OneInputStreamOperatorTestHarness<RowData, RowData> testHarness =
                createTestHarness(function);
        testHarness.open();

        // Process a DELETE for key "aaa" (no existing state, will be filtered out)
        testHarness.processElement(deleteRecord("aaa", 1L));

        // Process INSERTs for keys "bbb" and "ccc"
        testHarness.processElement(insertRecord("bbb", 2L));
        testHarness.processElement(insertRecord("ccc", 3L));

        // Close to trigger finishBundle
        testHarness.close();

        // Verify that keys "bbb" and "ccc" were processed correctly
        List<Object> expectedOutput = new ArrayList<>();
        expectedOutput.add(insertRecord("bbb", 2L));
        expectedOutput.add(insertRecord("ccc", 3L));

        assertor.assertOutputEqualsSorted(
                "Keys after retraction-only key should still be processed",
                expectedOutput,