Repository: apache/kafka
File: streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/EosIntegrationTest.java (Java)

Code before:
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C

            final List<KeyValue<Long, Long>> committedRecordsAfterRebalance = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size(),
                CONSUMER_GROUP_ID);

            final List<KeyValue<Long, Long>> expectedCommittedRecordsAfterRebalance = new ArrayList<>(
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size());
            expectedCommittedRecordsAfterRebalance.addAll(uncommittedDataBeforeStall);
            expectedCommittedRecordsAfterRebalance.addAll(dataToTriggerFirstRebalance);

            checkResultPerKey(
                committedRecordsAfterRebalance,
                expectedCommittedRecordsAfterRebalance,
                "The all committed records after rebalance do not match what expected");

            doStall = false;
            // After removing the stalling instance, we expect both remaining instances to see both instances.
            // It doesn't really matter what the assignment is, but we might as well also assert that they
            // both see both partitions assigned exactly once
            waitForCondition(
                () -> streams1.metadataForAllStreamsClients().size() == 2
                    && streams2.metadataForAllStreamsClients().size() == 2
                    && streams1.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2
                    && streams2.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2,
                MAX_WAIT_TIME_MS,
                () -> "Should have rebalanced.\n" +
                    "Streams1[" + streams1.metadataForAllStreamsClients() + "]\n" +
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            writeInputData(dataAfterSecondRebalance);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C

            final List<KeyValue<Long, Long>> allCommittedRecords = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                committedDataBeforeStall.size() + uncommittedDataBeforeStall.size()
                + dataToTriggerFirstRebalance.size() + dataAfterSecondRebalance.size(),
                CONSUMER_GROUP_ID + "_ALL");

            final int allCommittedRecordsAfterRecoverySize = committedDataBeforeStall.size() +
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size() + dataAfterSecondRebalance.size();

Review comment (suggestion):
The comment mentions 'both remaining instances' but the code only has one remaining instance after the stalling instance is removed. This is inconsistent with the actual implementation where only one instance remains active.
```suggestion
            // After removing the stalling instance, we expect the single remaining instance to see only itself as a member.
            // It doesn't really matter what the assignment is, but we might as well also assert that it
            // sees both partitions assigned exactly once
```

Code after:
                () -> "Should have rebalanced.\n" +
                    "Streams1[" + streams1.metadataForAllStreamsClients() + "]\n" +
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C

            final List<KeyValue<Long, Long>> committedRecordsAfterRebalance = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size(),
                CONSUMER_GROUP_ID);

            final List<KeyValue<Long, Long>> expectedCommittedRecordsAfterRebalance = new ArrayList<>(
                uncommittedDataBeforeStall.size() + dataToTriggerFirstRebalance.size());
            expectedCommittedRecordsAfterRebalance.addAll(uncommittedDataBeforeStall);
            expectedCommittedRecordsAfterRebalance.addAll(dataToTriggerFirstRebalance);

            checkResultPerKey(
                committedRecordsAfterRebalance,
                expectedCommittedRecordsAfterRebalance,
                "The all committed records after rebalance do not match what expected");

            LOG.info("Releasing Stall");
            doStall = false;
            // Once the stalling host rejoins the group, we expect both instances to see both instances.
            // It doesn't really matter what the assignment is, but we might as well also assert that they
            // both see both partitions assigned exactly once
            waitForCondition(
                () -> streams1.metadataForAllStreamsClients().size() == 2
                    && streams2.metadataForAllStreamsClients().size() == 2
                    && streams1.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2
                    && streams2.metadataForAllStreamsClients().stream().mapToLong(meta -> meta.topicPartitions().size()).sum() == 2,
                MAX_WAIT_TIME_MS,
                () -> "Should have rebalanced.\n" +
                    "Streams1[" + streams1.metadataForAllStreamsClients() + "]\n" +
                    "Streams2[" + streams2.metadataForAllStreamsClients() + "]");

            writeInputData(dataAfterSecondRebalance);

            // expected end state per output partition (C == COMMIT; A == ABORT; ---> indicate the changes):
            //
            // p-0: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C
            // p-1: ---> 10 rec + C  + 5 rec + C    + 5 rec + C   + 10 rec + C

            final List<KeyValue<Long, Long>> allCommittedRecords = readResult(
                SINGLE_PARTITION_OUTPUT_TOPIC,
                committedDataBeforeStall.size() + uncommittedDataBeforeStall.size()
                + dataToTriggerFirstRebalance.size() + dataAfterSecondRebalance.size(),
                CONSUMER_GROUP_ID + "_ALL");