feat: enhance batch matching with expected record keys

src/main/java/io/github/ktestify/io/kafka/AbstractKafkaConsumer.java

@@ -123,10 +123,10 @@ protected MatchContext buildMatchContext() {
                 .matchMethod(context.getMatchMethod())
                 .matchFilePaths(context.getMatchFilePaths())
                 .matchKey(context.getExpectedRecordKey())
+                .matchKeys(context.getExpectedRecordKeys())
                 .keyMatchStrategy(context.getKeyMatchStrategy())
                 .excludedFields(context.getExcludedFields())
                 .strictMatching(false)
-                .keyMatchStrategy(context.getKeyMatchStrategy())
                 .batchMatchStrategy(context.getBatchMatchStrategy())
                 .build();
     }

src/main/java/io/github/ktestify/io/kafka/AbstractKafkaProducer.java

@@ -145,11 +145,13 @@ protected String resolveSchema() throws IOException, RestClientException {
         } else {
             try {
                 int version = Integer.parseInt(context.getSchemaVersion());
-                logger.info("Resolving Avro schema for subject '{}' (version {}) from Schema Registry.", subject, version);
+                logger.info(
+                        "Resolving Avro schema for subject '{}' (version {}) from Schema Registry.", subject, version);
                 return client.getSchemaMetadata(subject, version).getSchema();
             } catch (NumberFormatException e) {
                 throw new ProducerException(
-                        "Invalid schema version '" + context.getSchemaVersion() + "' for subject '" + subject + "'. Expected an integer.",
+                        "Invalid schema version '" + context.getSchemaVersion() + "' for subject '" + subject
+                                + "'. Expected an integer.",
                         e);
             }
         }

src/main/java/io/github/ktestify/io/kafka/ConsumerContext.java

@@ -33,6 +33,7 @@ public final class ConsumerContext<K, V> {
     private final Map<String, String> properties;
     private final Consumer<K, V> consumer;
     private final String expectedRecordKey;
+    private final List<String> expectedRecordKeys;
     private final KeyMatchStrategy keyMatchStrategy;
     private final BatchMatchStrategy batchMatchStrategy;
     private final String matchMethod;
@@ -49,6 +50,7 @@ private ConsumerContext(
             Map<String, String> properties,
             Consumer<K, V> consumer,
             String expectedRecordKey,
+            List<String> expectedRecordKeys,
             KeyMatchStrategy keyMatchStrategy,
             BatchMatchStrategy batchMatchStrategy,
             String matchMethod,
@@ -63,6 +65,7 @@ private ConsumerContext(
         this.properties = properties;
         this.consumer = consumer;
         this.expectedRecordKey = expectedRecordKey;
+        this.expectedRecordKeys = expectedRecordKeys != null ? expectedRecordKeys : Collections.emptyList();
         this.keyMatchStrategy = keyMatchStrategy != null ? keyMatchStrategy : KeyMatchStrategy.EXACT;
         this.batchMatchStrategy = batchMatchStrategy != null ? batchMatchStrategy : BatchMatchStrategy.POSITIONAL;
         this.matchMethod = matchMethod;
@@ -93,6 +96,7 @@ public static final class Builder<K, V> {
         private Map<String, String> properties;
         private Consumer<K, V> consumer;
         private String expectedRecordKey;
+        private List<String> expectedRecordKeys;
         private KeyMatchStrategy keyMatchStrategy;
         private BatchMatchStrategy batchMatchStrategy;
         private String matchMethod;
@@ -124,6 +128,21 @@ public Builder<K, V> expectedRecordKey(String expectedRecordKey) {
             return this;
         }
 
+        /**
+         * Sets the expected record keys for batch consumers.
+         *
+         * <p>When non-empty, {@code KafkaRecordFetcher} pre-filters records against this list of expected keys, and
+         * batch matchers assert each record's key alongside its value according to {@link #batchMatchStrategy}.
+         *
+         * @param expectedRecordKeys the expected record keys, or {@code null} for an empty list
+         * @return this builder
+         * @since 1.2.6
+         */
+        public Builder<K, V> expectedRecordKeys(List<String> expectedRecordKeys) {
+            this.expectedRecordKeys = expectedRecordKeys;
+            return this;
+        }
+
         /**
          * Sets the strategy used to compare {@link #expectedRecordKey} against the actual record key during the
          * fetch-time pre-filter in {@code KafkaRecordFetcher.passesKeyFilter()}.
@@ -235,6 +254,7 @@ public ConsumerContext<K, V> build() {
                     validatedProps,
                     validatedConsumer,
                     expectedRecordKey,
+                    expectedRecordKeys,
                     keyMatchStrategy,
                     batchMatchStrategy,
                     matchMethod,

src/main/java/io/github/ktestify/io/kafka/KafkaRecordFetcher.java

@@ -307,15 +307,20 @@ private List<ConsumedRecord<V>> pollUntilRecordFound(long deadlineNanos) throws
         }
         List<ConsumedRecord<V>> accumulated = new ArrayList<>();
 
+        String keyFilterDesc = "<any>";
+        if (context.getExpectedRecordKeys() != null
+                && !context.getExpectedRecordKeys().isEmpty()) {
+            keyFilterDesc = context.getExpectedRecordKeys() + " (" + context.getKeyMatchStrategy() + ")";
+        } else if (context.getExpectedRecordKey() != null) {
+            keyFilterDesc = "'" + context.getExpectedRecordKey() + "' (" + context.getKeyMatchStrategy() + ")";
+        }
         log.info(
                 "Polling topic '{}' in {} mode (target={} record(s), timeout={}ms, expectedKey={}).",
                 namespacedTopic,
                 batchMode ? "BATCH" : "SINGLE",
                 targetSize,
                 readTimeoutMs,
-                context.getExpectedRecordKey() != null
-                        ? "'" + context.getExpectedRecordKey() + "' (" + context.getKeyMatchStrategy() + ")"
-                        : "<any>");
+                keyFilterDesc);
 
         while (System.nanoTime() < deadlineNanos) {
             ConsumerRecords<K, V> batch = kafkaConsumer.poll(Duration.ofMillis(pollIntervalMs));
@@ -412,7 +417,22 @@ private void release(ConsumedRecord<V> record) {
      * context / properties using the configured {@link io.github.ktestify.match.KeyMatchStrategy}.
      */
     private boolean passesKeyFilter(ConsumerRecord<K, V> record) {
-        // Context takes priority over properties map
+        String recordKey = record.key() != null ? record.key().toString() : null;
+
+        // 1. Batch expectedRecordKeys list (if non-empty)
+        List<String> expectedKeys = context.getExpectedRecordKeys();
+        if (expectedKeys != null && !expectedKeys.isEmpty()) {
+            for (String expectedKey : expectedKeys) {
+                if (context.getKeyMatchStrategy().matches(expectedKey, recordKey)) {
+                    log.info(MESSAGE_CONSUMER_RECORD_MATCHES_EXPECTED_KEY, expectedKey);
+                    return true;
+                }
+            }
+            log.debug(MESSAGE_CONSUMER_RECORD_DO_NOT_MATCHES_EXPECTED_KEY, recordKey, expectedKeys);
+            return false;
+        }
+
+        // 2. Single expectedRecordKey (context takes priority over properties map)
         String expectedKey = context.getExpectedRecordKey();
         if (expectedKey == null || expectedKey.isEmpty()) {
             expectedKey = properties.get("expectedRecordKey");
@@ -422,7 +442,6 @@ private boolean passesKeyFilter(ConsumerRecord<K, V> record) {
             return true; // No filter configured: accept any key
         }
 
-        String recordKey = record.key() != null ? record.key().toString() : null;
         if (context.getKeyMatchStrategy().matches(expectedKey, recordKey)) {
             log.info(MESSAGE_CONSUMER_RECORD_MATCHES_EXPECTED_KEY, expectedKey);
             return true;

src/main/java/io/github/ktestify/match/MatchContext.java

@@ -70,6 +70,19 @@ public class MatchContext {
     /** Optional key whose value must match a specific value (key-value matching mode). */
     String matchKey;
 
+    /**
+     * Expected record keys for batch matching.
+     *
+     * <p>When non-empty in batch matchers ({@code FileRecordMatcher}, {@code AvroFileRecordMatcher}), must have the
+     * same size as {@link #matchFilePaths}. In {@link BatchMatchStrategy#POSITIONAL} mode, record {@code i} must match
+     * both {@code matchKeys.get(i)} and {@code matchFilePaths.get(i)}; in {@link BatchMatchStrategy#LAZY} mode, each
+     * record claims the first unmatched {@code (key, file)} pair where both key and value match.
+     *
+     * @since 1.2.6
+     */
+    @Builder.Default
+    List<String> matchKeys = Collections.emptyList();
+
     /** Expected value for {@link #matchKey}. */
     String matchValue;
 

src/main/java/io/github/ktestify/match/impl/AvroFileRecordMatcher.java

@@ -16,10 +16,7 @@
 package io.github.ktestify.match.impl;
 
 import io.github.ktestify.exceptions.ComparisonException;
-import io.github.ktestify.match.BatchMatchStrategy;
-import io.github.ktestify.match.MatchContext;
-import io.github.ktestify.match.MatchResult;
-import io.github.ktestify.match.RecordMatcher;
+import io.github.ktestify.match.*;
 import io.github.ktestify.models.ConsumedRecord;
 import io.github.ktestify.utils.FileUtils;
 import io.github.ktestify.utils.serdes.AvroUtils;
@@ -51,8 +48,13 @@ public MatchResult match(List<ConsumedRecord<GenericRecord>> records, MatchConte
         if (filePaths == null || filePaths.isEmpty() || filePaths.stream().anyMatch(p -> p == null || p.isBlank())) {
             throw new ComparisonException("AvroFileRecordMatcher requires matchFilePath(s) to be set.");
         }
+        List<String> expectedKeys = context.getMatchKeys() != null ? context.getMatchKeys() : List.of();
+        if (!expectedKeys.isEmpty() && expectedKeys.size() != filePaths.size()) {
+            throw new ComparisonException("Expected record keys count (" + expectedKeys.size()
+                    + ") does not match expected file count (" + filePaths.size() + ").");
+        }
 
-        if (records.size() == 1 && filePaths.size() == 1) {
+        if (records.size() == 1 && filePaths.size() == 1 && expectedKeys.isEmpty()) {
             return matchSingle(records.getFirst(), filePaths.getFirst(), context);
         }
 
@@ -66,9 +68,9 @@ public MatchResult match(List<ConsumedRecord<GenericRecord>> records, MatchConte
                 : BatchMatchStrategy.POSITIONAL;
 
         if (strategy == BatchMatchStrategy.LAZY) {
-            return matchBatchLazy(records, filePaths, context);
+            return matchBatchLazy(records, filePaths, expectedKeys, context);
         }
-        return matchBatchPositional(records, filePaths, context);
+        return matchBatchPositional(records, filePaths, expectedKeys, context);
     }
 
     private MatchResult matchSingle(ConsumedRecord<GenericRecord> record, String filePath, MatchContext context) {
@@ -95,13 +97,37 @@ private MatchResult matchSingle(ConsumedRecord<GenericRecord> record, String fil
     }
 
     private MatchResult matchBatchPositional(
-            List<ConsumedRecord<GenericRecord>> records, List<String> filePaths, MatchContext context) {
+            List<ConsumedRecord<GenericRecord>> records,
+            List<String> filePaths,
+            List<String> expectedKeys,
+            MatchContext context) {
+        KeyMatchStrategy keyStrategy =
+                context.getKeyMatchStrategy() != null ? context.getKeyMatchStrategy() : KeyMatchStrategy.EXACT;
         String lastExpected = "";
         String lastActual = "";
         for (int i = 0; i < records.size(); i++) {
             String filePath = filePaths.get(i);
             String expectedValue = FileUtils.getFileContent(FileUtils.getFile(filePath));
             GenericRecord value = records.get(i).getValue();
+            String actualKey = records.get(i).getKey();
+
+            if (!expectedKeys.isEmpty()) {
+                String expectedKey = expectedKeys.get(i);
+                if (!keyStrategy.matches(expectedKey, actualKey)) {
+                    log.error(
+                            "Avro record at index {} key mismatch: expected '{}' ({}), actual '{}'.",
+                            i,
+                            expectedKey,
+                            keyStrategy,
+                            actualKey);
+                    return MatchResult.fail(
+                            "Avro record at index " + i + " key does not match, expected: '" + expectedKey
+                                    + "', actual: '" + actualKey + "', strategy: " + keyStrategy + ".",
+                            expectedKey,
+                            actualKey != null ? actualKey : "<null>");
+                }
+            }
+
             if (value == null) {
                 log.error("Avro record at index {} has null value (tombstone), expected content of '{}'.", i, filePath);
                 return MatchResult.nullValue(expectedValue);
@@ -119,15 +145,28 @@ private MatchResult matchBatchPositional(
                         expectedValue,
                         actualValue);
             }
-            log.info("Avro record at index {} matches expected file '{}'.", i, filePath);
+            if (!expectedKeys.isEmpty()) {
+                log.info(
+                        "Avro record at index {} matches expected key '{}' and file '{}'.",
+                        i,
+                        expectedKeys.get(i),
+                        filePath);
+            } else {
+                log.info("Avro record at index {} matches expected file '{}'.", i, filePath);
+            }
             lastExpected = expectedValue;
             lastActual = actualValue;
         }
         return MatchResult.pass(lastExpected, lastActual);
     }
 
     private MatchResult matchBatchLazy(
-            List<ConsumedRecord<GenericRecord>> records, List<String> filePaths, MatchContext context) {
+            List<ConsumedRecord<GenericRecord>> records,
+            List<String> filePaths,
+            List<String> expectedKeys,
+            MatchContext context) {
+        KeyMatchStrategy keyStrategy =
+                context.getKeyMatchStrategy() != null ? context.getKeyMatchStrategy() : KeyMatchStrategy.EXACT;
         List<String> expectedContents = filePaths.stream()
                 .map(path -> FileUtils.getFileContent(FileUtils.getFile(path)))
                 .toList();
@@ -137,6 +176,7 @@ private MatchResult matchBatchLazy(
 
         for (int i = 0; i < records.size(); i++) {
             GenericRecord value = records.get(i).getValue();
+            String actualKey = records.get(i).getKey();
             if (value == null) {
                 log.error("Avro record at index {} has null value (tombstone).", i);
                 return MatchResult.nullValue(expectedContents.getFirst());
@@ -145,31 +185,45 @@ private MatchResult matchBatchLazy(
 
             int matchedIndex = -1;
             for (int j = 0; j < filePaths.size(); j++) {
-                if (!matchedFiles[j] && matchesExpectedContentQuietly(expectedContents.get(j), actualValue, context)) {
+                if (!matchedFiles[j]
+                        && (expectedKeys.isEmpty() || keyStrategy.matches(expectedKeys.get(j), actualKey))
+                        && matchesExpectedContentQuietly(expectedContents.get(j), actualValue, context)) {
                     matchedIndex = j;
                     break;
                 }
             }
 
             if (matchedIndex >= 0) {
                 matchedFiles[matchedIndex] = true;
-                log.info("Avro record at index {} matched expected file '{}'.", i, filePaths.get(matchedIndex));
+                if (!expectedKeys.isEmpty()) {
+                    log.info(
+                            "Avro record at index {} matched expected key '{}' and file '{}'.",
+                            i,
+                            expectedKeys.get(matchedIndex),
+                            filePaths.get(matchedIndex));
+                } else {
+                    log.info("Avro record at index {} matched expected file '{}'.", i, filePaths.get(matchedIndex));
+                }
                 lastExpected = expectedContents.get(matchedIndex);
                 lastActual = actualValue;
             } else {
                 List<String> unmatchedFiles = new ArrayList<>();
                 String firstUnmatchedExpected = "";
                 for (int j = 0; j < filePaths.size(); j++) {
                     if (!matchedFiles[j]) {
-                        unmatchedFiles.add(filePaths.get(j));
+                        unmatchedFiles.add(
+                                !expectedKeys.isEmpty()
+                                        ? expectedKeys.get(j) + " -> " + filePaths.get(j)
+                                        : filePaths.get(j));
                         if (firstUnmatchedExpected.isEmpty()) {
                             firstUnmatchedExpected = expectedContents.get(j);
                         }
                     }
                 }
                 log.error(
-                        "Avro record at index {} did not match any remaining expected file. Unmatched files: {}\nActual record:\n{}",
+                        "Avro record at index {} (key='{}') did not match any remaining expected file. Unmatched files: {}\nActual record:\n{}",
                         i,
+                        actualKey,
                         unmatchedFiles,
                         actualValue);
                 return MatchResult.fail(