Repository: apache/kafka
File: clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerWithLegacyMessageFormatIntegrationTest.java (Java)

Code before:
    public void testOffsetsForTimesClassicConsumer() throws InterruptedException {
        testOffsetsForTimes(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testOffsetsForTimesWithAsyncConsumer() throws InterruptedException {
        testOffsetsForTimes(GroupProtocol.CONSUMER);
    }

    public void testOffsetsForTimes(GroupProtocol groupProtocol) throws InterruptedException {
        setupTopics();

        try (Consumer<Object, Object> consumer = cluster.consumer(Map.of(
                GROUP_PROTOCOL_CONFIG, groupProtocol.name.toLowerCase(Locale.ROOT)))
        ) {
            // Test negative target time
            assertThrows(IllegalArgumentException.class, () ->
                    consumer.offsetsForTimes(Map.of(t1p0, -1L)));

            Map<TopicPartition, Long> timestampsToSearch = new HashMap<>();
            timestampsToSearch.put(t1p0, 0L);
            timestampsToSearch.put(t1p1, 20L);
            timestampsToSearch.put(t2p0, 40L);
            timestampsToSearch.put(t2p1, 60L);
            timestampsToSearch.put(t3p0, 80L);
            timestampsToSearch.put(t3p1, 100L);

            Map<TopicPartition, OffsetAndTimestamp> timestampOffsets = consumer.offsetsForTimes(timestampsToSearch);

            OffsetAndTimestamp timestampTopic1P0 = timestampOffsets.get(t1p0);
            assertEquals(0, timestampTopic1P0.offset());
            assertEquals(0, timestampTopic1P0.timestamp());
            assertEquals(Optional.of(0), timestampTopic1P0.leaderEpoch());

            OffsetAndTimestamp timestampTopic1P1 = timestampOffsets.get(t1p1);
            assertEquals(20, timestampTopic1P1.offset());
            assertEquals(20, timestampTopic1P1.timestamp());
            assertEquals(Optional.of(0), timestampTopic1P1.leaderEpoch());

            // v0 message format doesn't have timestamp
            OffsetAndTimestamp timestampTopic2P0 = timestampOffsets.get(t2p0);
            assertNull(timestampTopic2P0);

            OffsetAndTimestamp timestampTopic2P1 = timestampOffsets.get(t2p1);
            assertNull(timestampTopic2P1);

            // v1 message format doesn't have leader epoch
            OffsetAndTimestamp timestampTopic3P0 = timestampOffsets.get(t3p0);
            assertEquals(80, timestampTopic3P0.offset());
            assertEquals(80, timestampTopic3P0.timestamp());
            assertEquals(Optional.empty(), timestampTopic3P0.leaderEpoch());

Review comment (suggestion):
```suggestion
            Map<TopicPartition, Long> timestampsToSearch = Map.of(
                t1p0, 0L,
                t1p1, 20L,
                t2p0, 40L,
                t2p1, 60L,
                t3p0, 80L,
                t3p1, 100L
            );
```

Code after:
                    t2p1, 60L,
                    t3p0, 80L,
                    t3p1, 100L
            );

            Map<TopicPartition, OffsetAndTimestamp> timestampOffsets = consumer.offsetsForTimes(timestampsToSearch);

            OffsetAndTimestamp timestampTopic1P0 = timestampOffsets.get(t1p0);
            assertEquals(0, timestampTopic1P0.offset());
            assertEquals(0, timestampTopic1P0.timestamp());
            assertEquals(Optional.of(0), timestampTopic1P0.leaderEpoch());

            OffsetAndTimestamp timestampTopic1P1 = timestampOffsets.get(t1p1);
            assertEquals(20, timestampTopic1P1.offset());
            assertEquals(20, timestampTopic1P1.timestamp());
            assertEquals(Optional.of(0), timestampTopic1P1.leaderEpoch());

            OffsetAndTimestamp timestampTopic2P0 = timestampOffsets.get(t2p0);
            assertNull(timestampTopic2P0, "v0 message format shouldn't have timestamp");

            OffsetAndTimestamp timestampTopic2P1 = timestampOffsets.get(t2p1);
            assertNull(timestampTopic2P1);

            OffsetAndTimestamp timestampTopic3P0 = timestampOffsets.get(t3p0);
            assertEquals(80, timestampTopic3P0.offset());
            assertEquals(80, timestampTopic3P0.timestamp());
            assertEquals(Optional.empty(), timestampTopic3P0.leaderEpoch());

            assertNull(timestampOffsets.get(t3p1), "v1 message format doesn't have leader epoch");
        }

    }

    @ClusterTest
    public void testEarliestOrLatestOffsetsWithClassicConsumer() {
        testEarliestOrLatestOffsets(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testEarliestOrLatestOffsetsWithAsyncConsumer() {
        testEarliestOrLatestOffsets(GroupProtocol.CONSUMER);
    }

    public void testEarliestOrLatestOffsets(GroupProtocol groupProtocol) {
        Set<TopicPartition> partitions = Set.of(t1p0, t1p1, t2p0, t2p1, t3p0, t3p1);

        try (Consumer<Object, Object> consumer = cluster.consumer(Map.of(
                GROUP_PROTOCOL_CONFIG, groupProtocol.name.toLowerCase(Locale.ROOT)))
        ) {
            Map<TopicPartition, Long> earliests = consumer.beginningOffsets(partitions);
            assertEquals(0L, earliests.get(t1p0));