Generalize version check assertion

Author: parkertimmins

server/src/test/java/org/elasticsearch/index/codec/tsdb/TsdbDocValueBwcTests.java

@@ -50,6 +50,7 @@
 import java.util.Arrays;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.IntSupplier;
 
 import static org.elasticsearch.index.codec.tsdb.es819.ES819TSDBDocValuesFormatTests.TestES819TSDBDocValuesFormatVersion0;
@@ -67,7 +68,12 @@ public void testMixedIndex() throws Exception {
     public void testMixedIndexDocValueVersion0ToCurrent() throws Exception {
         var oldCodec = TestUtil.alwaysDocValuesFormat(new TestES819TSDBDocValuesFormatVersion0());
         var newCodec = TestUtil.alwaysDocValuesFormat(new ES819TSDBDocValuesFormat());
-        testMixedIndex(oldCodec, newCodec);
+        testMixedIndex(
+            oldCodec,
+            newCodec,
+            this::assertVersion819,
+            this::assertVersion819
+        );
     }
 
     public void testMixedIndex816To900Lucene101() throws Exception {
@@ -92,7 +98,34 @@ public DocValuesFormat getDocValuesFormatForField(String field) {
         testMixedIndex(oldCodec, newCodec);
     }
 
-    void testMixedIndex(Codec oldCodec, Codec newCodec) throws IOException, NoSuchFieldException, IllegalAccessException,
+    void assertFieldInfoDocValuesFormat(DirectoryReader reader, String expectedSuffix, String expectedFormat) throws IOException, NoSuchFieldException, IllegalAccessException {
+        // Assert per field format field info attributes:
+        // (XPerFieldDocValuesFormat must produce the same attributes as PerFieldDocValuesFormat for BWC.
+        // Otherwise, doc values fields may disappear)
+        for (var leaf : reader.leaves()) {
+            for (var fieldInfo : leaf.reader().getFieldInfos()) {
+                assertThat(fieldInfo.attributes(), Matchers.aMapWithSize(2));
+                assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.suffix", expectedSuffix));
+                assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.format", expectedFormat));
+            }
+        }
+    }
+
+    void assertVersion87(DirectoryReader reader) throws IOException, NoSuchFieldException, IllegalAccessException {
+        assert87DocValuesFormatVersion(reader);
+        assertFieldInfoDocValuesFormat(reader, "0", "ES87TSDB");
+    }
+
+    void assertVersion819(DirectoryReader reader) throws IOException, NoSuchFieldException, ClassNotFoundException, IllegalAccessException {
+        assert819DocValuesFormatVersion(reader);
+        assertFieldInfoDocValuesFormat(reader, "0", "ES819TSDB");
+    }
+
+    void testMixedIndex(Codec oldCodec, Codec newCodec) throws IOException, NoSuchFieldException, IllegalAccessException, ClassNotFoundException {
+        testMixedIndex(oldCodec, newCodec, this::assertVersion87, this::assertVersion819);
+    }
+
+    void testMixedIndex(Codec oldCodec, Codec newCodec, VersionAssert assertOldVersion, VersionAssert assertNewVersion) throws IOException, NoSuchFieldException, IllegalAccessException,
         ClassNotFoundException {
         String timestampField = "@timestamp";
         String hostnameField = "host.name";
@@ -144,17 +177,7 @@ void testMixedIndex(Codec oldCodec, Codec newCodec) throws IOException, NoSuchFi
             }
             // Check documents before force merge:
             try (var reader = DirectoryReader.open(dir)) {
-                assertOldDocValuesFormatVersion(reader);
-                // Assert per field format field info attributes:
-                // (XPerFieldDocValuesFormat must produce the same attributes as PerFieldDocValuesFormat for BWC.
-                // Otherwise, doc values fields may disappear)
-                for (var leaf : reader.leaves()) {
-                    for (var fieldInfo : leaf.reader().getFieldInfos()) {
-                        assertThat(fieldInfo.attributes(), Matchers.aMapWithSize(2));
-                        assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.suffix", "0"));
-                        assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.format", "ES87TSDB"));
-                    }
-                }
+                assertOldVersion.run(reader);
 
                 var hostNameDV = MultiDocValues.getSortedValues(reader, hostnameField);
                 assertNotNull(hostNameDV);
@@ -213,17 +236,9 @@ void testMixedIndex(Codec oldCodec, Codec newCodec) throws IOException, NoSuchFi
                 try (var reader = DirectoryReader.open(iw)) {
                     assertEquals(1, reader.leaves().size());
                     assertEquals(numDocs, reader.maxDoc());
-                    assertNewDocValuesFormatVersion(reader);
-                    var leaf = reader.leaves().get(0).reader();
-                    // Assert per field format field info attributes:
-                    // (XPerFieldDocValuesFormat must produce the same attributes as PerFieldDocValuesFormat for BWC.
-                    // Otherwise, doc values fields may disappear)
-                    for (var fieldInfo : leaf.getFieldInfos()) {
-                        assertThat(fieldInfo.attributes(), Matchers.aMapWithSize(2));
-                        assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.suffix", "0"));
-                        assertThat(fieldInfo.attributes(), Matchers.hasEntry("PerFieldDocValuesFormat.format", "ES819TSDB"));
-                    }
+                    assertNewVersion.run(reader);
 
+                    var leaf = reader.leaves().get(0).reader();
                     var hostNameDV = leaf.getSortedDocValues(hostnameField);
                     assertNotNull(hostNameDV);
                     var timestampDV = DocValues.unwrapSingleton(leaf.getSortedNumericDocValues(timestampField));
@@ -393,7 +408,7 @@ private IndexWriterConfig getTimeSeriesIndexWriterConfig(String hostnameField, S
 
     // A hacky way to figure out whether doc values format is written in what version. Need to use reflection, because
     // PerFieldDocValuesFormat hides the doc values formats it wraps.
-    private void assertOldDocValuesFormatVersion(DirectoryReader reader) throws NoSuchFieldException, IllegalAccessException, IOException {
+    private void assert87DocValuesFormatVersion(DirectoryReader reader) throws NoSuchFieldException, IllegalAccessException, IOException {
         if (System.getSecurityManager() != null) {
             // With jvm version 24 entitlements are used and security manager is nog longer used.
             // Making this assertion work with security manager requires granting the entire test codebase privileges to use
@@ -414,7 +429,7 @@ private void assertOldDocValuesFormatVersion(DirectoryReader reader) throws NoSu
         }
     }
 
-    private void assertNewDocValuesFormatVersion(DirectoryReader reader) throws NoSuchFieldException, IllegalAccessException, IOException,
+    private void assert819DocValuesFormatVersion(DirectoryReader reader) throws NoSuchFieldException, IllegalAccessException, IOException,
         ClassNotFoundException {
 
         for (var leafReaderContext : reader.leaves()) {
@@ -459,4 +474,7 @@ private static Field getFormatsFieldFromPerFieldFieldsReader(Class<?> c) throws
         return field;
     }
 
+    interface VersionAssert {
+        void run(DirectoryReader reader) throws IOException, NoSuchFieldException, IllegalAccessException, ClassNotFoundException;
+    }
 }