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

Code before:
        cluster.createTopic(topic1, 2, (short) 1);
        createTopicWithAssignment(topic2, Map.of(0, List.of(0), 1, List.of(1)));
        createTopicWithAssignment(topic3, Map.of(0, List.of(0), 1, List.of(1)));

        // v2 message format for topic1
        ClientsTestUtils.sendRecords(cluster, t1p0, 100, 0);
        ClientsTestUtils.sendRecords(cluster, t1p1, 100, 0);
        // v0 message format for topic2
        appendLegacyRecords(100, t2p0, 0, RecordBatch.MAGIC_VALUE_V0);
        appendLegacyRecords(100, t2p1, 1, RecordBatch.MAGIC_VALUE_V0);
        // v1 message format for topic3
        appendLegacyRecords(100, t3p0, 0, RecordBatch.MAGIC_VALUE_V1);
        appendLegacyRecords(100, t3p1, 1, RecordBatch.MAGIC_VALUE_V1);
    }

    @ClusterTest
    public void testOffsetsForTimesWithClassicConsumer() throws InterruptedException {
        testOffsetsForTimes(GroupProtocol.CLASSIC);
    }

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

    public void testOffsetsForTimes(GroupProtocol groupProtocol) throws InterruptedException {
        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 = Map.of(
                    t1p0, 0L,
                    t1p1, 20L,
                    t2p0, 40L,
                    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());

Review comment (suggestion):
```suggestion
    public void testOffsetsForTimes(GroupProtocol groupProtocol) {
```

Code after:
        cluster.createTopic(topic1, 2, (short) 1);
        createTopicWithAssignment(topic2, Map.of(0, List.of(0), 1, List.of(1)));
        createTopicWithAssignment(topic3, Map.of(0, List.of(0), 1, List.of(1)));

        // v2 message format for topic1
        ClientsTestUtils.sendRecords(cluster, t1p0, 100, 0);
        ClientsTestUtils.sendRecords(cluster, t1p1, 100, 0);
        // v0 message format for topic2
        appendLegacyRecords(100, t2p0, 0, RecordBatch.MAGIC_VALUE_V0);
        appendLegacyRecords(100, t2p1, 1, RecordBatch.MAGIC_VALUE_V0);
        // v1 message format for topic3
        appendLegacyRecords(100, t3p0, 0, RecordBatch.MAGIC_VALUE_V1);
        appendLegacyRecords(100, t3p1, 1, RecordBatch.MAGIC_VALUE_V1);
    }

    @ClusterTest
    public void testOffsetsForTimesWithClassicConsumer() {
        testOffsetsForTimes(GroupProtocol.CLASSIC);
    }

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

    public void testOffsetsForTimes(GroupProtocol groupProtocol) {
        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 = Map.of(
                    t1p0, 0L,
                    t1p1, 20L,
                    t2p0, 40L,
                    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());