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

Code before:
        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
        sendRecords(100, t1p0, 0);
        sendRecords(100, t1p1, 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 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());

Review comment (suggestion):
Use `@BeforeEach` instead of manually separating the setup logic.

Code after:
        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());
            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);