From b02b7d8481e212c1a33348d63227e54524cab4a7 Mon Sep 17 00:00:00 2001
From: xicm <xicm@asiainfo.com>
Date: Wed, 8 Jun 2022 10:28:54 +0800
Subject: [PATCH] HUDI-3682 testReaderFilterRowKeys fails in
 TestHoodieOrcReaderWriter

---
 .../io/storage/TestHoodieOrcReaderWriter.java  | 10 +++++-----
 .../io/storage/TestHoodieReaderWriterBase.java | 18 +++++++++++++++++-
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
index 282f102697d78..817fc25a5d639 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieOrcReaderWriter.java
@@ -34,6 +34,7 @@
 import org.mockito.Mockito;
 
 import java.io.IOException;
+import java.util.function.Supplier;
 
 import static org.apache.hudi.avro.HoodieAvroWriteSupport.HOODIE_AVRO_BLOOM_FILTER_METADATA_KEY;
 import static org.apache.hudi.avro.HoodieAvroWriteSupport.HOODIE_MAX_RECORD_KEY_FOOTER;
@@ -41,6 +42,7 @@
 import static org.apache.hudi.io.storage.HoodieOrcConfig.AVRO_SCHEMA_METADATA_KEY;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.when;
 
 public class TestHoodieOrcReaderWriter extends TestHoodieReaderWriterBase {
 
@@ -59,6 +61,9 @@ protected HoodieFileWriter<GenericRecord> createWriter(
     int maxFileSize = Integer.parseInt(HoodieStorageConfig.ORC_FILE_MAX_SIZE.defaultValue());
     HoodieOrcConfig config = new HoodieOrcConfig(conf, CompressionKind.ZLIB, orcStripSize, orcBlockSize, maxFileSize, filter);
     TaskContextSupplier mockTaskContextSupplier = Mockito.mock(TaskContextSupplier.class);
+    Supplier<Integer> partitionSupplier = Mockito.mock(Supplier.class);
+    when(mockTaskContextSupplier.getPartitionIdSupplier()).thenReturn(partitionSupplier);
+    when(partitionSupplier.get()).thenReturn(10);
     String instantTime = "000";
     return new HoodieOrcWriter<>(instantTime, getFilePath(), config, avroSchema, mockTaskContextSupplier);
   }
@@ -92,9 +97,4 @@ protected void verifySchema(Configuration conf, String schemaPath) throws IOExce
           orcReader.getSchema().toString());
     }
   }
-
-  @Override
-  public void testReaderFilterRowKeys() {
-    // TODO(HUDI-3682): fix filterRowKeys test for ORC due to a bug in ORC logic
-  }
 }
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
index 4617eb93a66e7..902f42e38f32b 100644
--- a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/io/storage/TestHoodieReaderWriterBase.java
@@ -20,6 +20,7 @@
 package org.apache.hudi.io.storage;
 
 import org.apache.hudi.common.bloom.BloomFilter;
+import org.apache.hudi.common.model.HoodieKey;
 
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericData;
@@ -157,7 +158,7 @@ public void testWriteReadWithEvolvedSchema() throws Exception {
 
   @Test
   public void testReaderFilterRowKeys() throws Exception {
-    writeFileWithSimpleSchema();
+    writeFileWithSchemaWithMeta();
     Configuration conf = new Configuration();
     verifyMetadata(conf);
     verifyFilterRowKeys(createReader(conf));
@@ -177,6 +178,21 @@ protected void writeFileWithSimpleSchema() throws Exception {
     writer.close();
   }
 
+  protected void writeFileWithSchemaWithMeta() throws Exception {
+    Schema avroSchema = getSchemaFromResource(TestHoodieReaderWriterBase.class, "/exampleSchemaWithMetaFields.avsc");
+    HoodieFileWriter<GenericRecord> writer = createWriter(avroSchema, true);
+    for (int i = 0; i < NUM_RECORDS; i++) {
+      GenericRecord record = new GenericData.Record(avroSchema);
+      String key = "key" + String.format("%02d", i);
+      record.put("_row_key", key);
+      record.put("time", Integer.toString(i));
+      record.put("number", i);
+      writer.writeAvroWithMetadata(new HoodieKey((String) record.get("_row_key"),
+          Integer.toString((Integer) record.get("number"))), record);
+    }
+    writer.close();
+  }
+
   protected void verifySimpleRecords(Iterator<GenericRecord> iterator) {
     int index = 0;
     while (iterator.hasNext()) {