From 2e82d0aa7b114000bcf58f9485377eb5705ba3b1 Mon Sep 17 00:00:00 2001
From: Ketan Verma <ketan9495@gmail.com>
Date: Mon, 18 Dec 2023 16:13:21 +0530
Subject: [PATCH] Improve execution of Roundable unit-tests

Signed-off-by: Ketan Verma <ketan9495@gmail.com>
---
 libs/common/build.gradle                      |  8 ---
 .../common/round/RoundableTests.java          | 68 +++++++++++++------
 2 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/libs/common/build.gradle b/libs/common/build.gradle
index 5dd25fde56a0f..704fe1b3d63da 100644
--- a/libs/common/build.gradle
+++ b/libs/common/build.gradle
@@ -91,14 +91,6 @@ if (BuildParams.runtimeJavaVersion >= JavaVersion.VERSION_20) {
     classpath -= sourceSets.main.output
   }
 
-  tasks.register('roundableSimdTest', Test) {
-    group 'verification'
-    include '**/RoundableTests.class'
-    systemProperty 'opensearch.experimental.feature.simd.rounding.enabled', 'true'
-  }
-
-  check.dependsOn(roundableSimdTest)
-
   forbiddenApisJava20 {
     failOnMissingClasses = false
     ignoreSignaturesOfMissingClasses = true
diff --git a/libs/common/src/test/java/org/opensearch/common/round/RoundableTests.java b/libs/common/src/test/java/org/opensearch/common/round/RoundableTests.java
index 1c0a36fe28ee8..3873a599ad8c7 100644
--- a/libs/common/src/test/java/org/opensearch/common/round/RoundableTests.java
+++ b/libs/common/src/test/java/org/opensearch/common/round/RoundableTests.java
@@ -8,36 +8,61 @@
 
 package org.opensearch.common.round;
 
+import org.opensearch.common.SuppressForbidden;
 import org.opensearch.test.OpenSearchTestCase;
 
-import java.util.List;
+import java.lang.invoke.MethodHandles;
+import java.lang.reflect.InvocationTargetException;
 
 public class RoundableTests extends OpenSearchTestCase {
 
-    public void testRoundingEmptyArray() {
-        Throwable throwable = assertThrows(IllegalArgumentException.class, () -> RoundableFactory.create(new long[0], 0));
-        assertEquals("at least one value must be present", throwable.getMessage());
+    public void testBidirectionalLinearSearcher() {
+        assertRounding(BidirectionalLinearSearcher::new);
     }
 
-    public void testRoundingSmallArray() {
-        int size = randomIntBetween(1, 64);
-        long[] values = randomArrayOfSortedValues(size);
-        Roundable roundable = RoundableFactory.create(values, size);
-
-        assertEquals("BidirectionalLinearSearcher", roundable.getClass().getSimpleName());
-        assertRounding(roundable, values, size);
+    public void testBinarySearcher() {
+        assertRounding(BinarySearcher::new);
     }
 
-    public void testRoundingLargeArray() {
-        int size = randomIntBetween(65, 256);
-        long[] values = randomArrayOfSortedValues(size);
-        Roundable roundable = RoundableFactory.create(values, size);
+    @SuppressForbidden(reason = "Reflective construction of BtreeSearcher since it's not supported below Java 20")
+    public void testBtreeSearcher() {
+        RoundableSupplier supplier;
+
+        try {
+            Class<?> clz = MethodHandles.lookup().findClass("org.opensearch.common.round.BtreeSearcher");
+            supplier = (values, size) -> {
+                try {
+                    return (Roundable) clz.getDeclaredConstructor(long[].class, int.class).newInstance(values, size);
+                } catch (InvocationTargetException e) {
+                    // Failed to instantiate the class. Unwrap if the nested exception is already a runtime exception,
+                    // say due to an IllegalArgumentException due to bad constructor arguments.
+                    if (e.getCause() instanceof RuntimeException) {
+                        throw (RuntimeException) e.getCause();
+                    } else {
+                        throw new RuntimeException(e);
+                    }
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            };
+        } catch (ClassNotFoundException e) {
+            assumeTrue("BtreeSearcher is not supported below Java 20", false);
+            return;
+        } catch (IllegalAccessException e) {
+            throw new RuntimeException(e);
+        }
 
-        assertTrue(List.of("BtreeSearcher", "BinarySearcher").contains(roundable.getClass().getSimpleName()));
-        assertRounding(roundable, values, size);
+        assertRounding(supplier);
     }
 
-    private void assertRounding(Roundable roundable, long[] values, int size) {
+    private void assertRounding(RoundableSupplier supplier) {
+        Throwable throwable = assertThrows(IllegalArgumentException.class, () -> supplier.get(new long[0], 0));
+        assertEquals("at least one value must be present", throwable.getMessage());
+
+        int size = randomIntBetween(1, 256);
+        long[] values = randomArrayOfSortedValues(size);
+        Roundable roundable = supplier.get(values, size);
+
         for (int i = 0; i < 100000; i++) {
             // Index of the expected round-down point.
             int idx = randomIntBetween(0, size - 1);
@@ -55,7 +80,7 @@ private void assertRounding(Roundable roundable, long[] values, int size) {
             assertEquals(expected, roundable.floor(key));
         }
 
-        Throwable throwable = assertThrows(AssertionError.class, () -> roundable.floor(values[0] - 1));
+        throwable = assertThrows(AssertionError.class, () -> roundable.floor(values[0] - 1));
         assertEquals("key must be greater than or equal to " + values[0], throwable.getMessage());
     }
 
@@ -69,4 +94,9 @@ private static long[] randomArrayOfSortedValues(int size) {
 
         return values;
     }
+
+    @FunctionalInterface
+    private interface RoundableSupplier {
+        Roundable get(long[] values, int size);
+    }
 }