test: Refactor unified query profiling tests with shared test base

Extract common test schema and utilities into UnifiedQueryTestBase to
reduce duplication across unified query test classes. Profiling tests
now extend this base and override buildContext() to enable profiling.

Change all timing assertions from > 0 to >= 0 to prevent flakiness.
Millisecond-resolution timers can return 0 for fast operations when
System.currentTimeMillis() samples the same tick before and after.
The >= 0 assertion still validates the metric was recorded (phase key
exists in the profile map) without depending on wall-clock granularity.

Signed-off-by: Chen Dai <daichen@amazon.com>

Author: dai-chen

api/src/test/java/org/opensearch/sql/api/UnifiedQueryProfilingTest.java

@@ -7,7 +7,6 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -21,143 +20,99 @@
 import org.junit.Test;
 import org.opensearch.sql.api.compiler.UnifiedQueryCompiler;
 import org.opensearch.sql.api.transpiler.UnifiedQueryTranspiler;
-import org.opensearch.sql.executor.QueryType;
 import org.opensearch.sql.monitor.profile.QueryProfile;
 import org.opensearch.sql.monitor.profile.QueryProfiling;
 
 /** Tests for profiling across unified query components and the measure() API. */
 public class UnifiedQueryProfilingTest extends UnifiedQueryTestBase {
 
-  @Test
-  public void testProfilingEnabled() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      assertTrue(QueryProfiling.current().isEnabled());
-    }
+  @Override
+  protected UnifiedQueryContext.Builder buildContext() {
+    return super.buildContext().profiling(true);
   }
 
   @Test
-  public void testProfilingDisabledByDefault() throws Exception {
-    try (UnifiedQueryContext ctx =
-        UnifiedQueryContext.builder()
-            .language(QueryType.PPL)
-            .catalog("opensearch", testSchema)
-            .defaultNamespace("opensearch")
-            .build()) {
-      assertFalse(QueryProfiling.current().isEnabled());
-    }
+  public void testProfilingEnabled() {
+    assertTrue(QueryProfiling.current().isEnabled());
   }
 
   @Test
-  public void testProfilingExplicitlyDisabled() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(false)) {
+  public void testProfilingDisabledByDefault() throws Exception {
+    try (UnifiedQueryContext ctx = super.buildContext().build()) {
       assertFalse(QueryProfiling.current().isEnabled());
     }
   }
 
   @Test
   public void testProfilingClearedAfterClose() throws Exception {
-    UnifiedQueryContext ctx = buildContext(true);
     assertTrue(QueryProfiling.current().isEnabled());
-    ctx.close();
+    context.close();
     assertFalse(QueryProfiling.current().isEnabled());
   }
 
   @Test
   public void testGetProfileReturnsNullWhenDisabled() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(false)) {
+    try (UnifiedQueryContext ctx = super.buildContext().build()) {
       assertNull(ctx.getProfile());
     }
   }
 
   @Test
-  public void testPlannerAutoProfilesAnalyzePhase() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      new UnifiedQueryPlanner(ctx).plan("source = opensearch.employees");
-      QueryProfile profile = ctx.getProfile();
-      assertNotNull(profile);
-      assertTrue(profile.getPhases().get("analyze").getTimeMillis() > 0);
-    }
+  public void testPlannerAutoProfilesAnalyzePhase() {
+    planner.plan("source = catalog.employees");
+    assertTrue(context.getProfile().getPhases().get("analyze").getTimeMillis() >= 0);
   }
 
   @Test
-  public void testCompilerAutoProfilesOptimizePhase() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      RelNode plan = new UnifiedQueryPlanner(ctx).plan("source = opensearch.employees");
-      PreparedStatement stmt = new UnifiedQueryCompiler(ctx).compile(plan);
-      assertNotNull(stmt);
-      QueryProfile profile = ctx.getProfile();
-      assertNotNull(profile);
-      assertTrue(profile.getPhases().get("optimize").getTimeMillis() > 0);
-    }
+  public void testCompilerAutoProfilesOptimizePhase() {
+    RelNode plan = planner.plan("source = catalog.employees");
+    new UnifiedQueryCompiler(context).compile(plan);
+    assertTrue(context.getProfile().getPhases().get("optimize").getTimeMillis() >= 0);
   }
 
   @Test
-  public void testTranspilerAutoProfilesTranspilePhase() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      RelNode plan = new UnifiedQueryPlanner(ctx).plan("source = opensearch.employees");
-      UnifiedQueryTranspiler transpiler = new UnifiedQueryTranspiler(ctx, SparkSqlDialect.DEFAULT);
-      assertNotNull(transpiler.toSql(plan));
-      QueryProfile profile = ctx.getProfile();
-      assertNotNull(profile);
-      assertTrue(profile.getPhases().get("transpile").getTimeMillis() > 0);
-    }
+  public void testTranspilerAutoProfilesTranspilePhase() {
+    RelNode plan = planner.plan("source = catalog.employees");
+    new UnifiedQueryTranspiler(context, SparkSqlDialect.DEFAULT).toSql(plan);
+    assertTrue(context.getProfile().getPhases().get("transpile").getTimeMillis() >= 0);
   }
 
   @Test
   public void testMeasureRecordsMetric() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      assertEquals("done", ctx.measure(EXECUTE, () -> "done"));
-      QueryProfile profile = ctx.getProfile();
-      assertNotNull(profile);
-      assertTrue(profile.getPhases().get("execute").getTimeMillis() >= 0);
-    }
+    assertEquals("done", context.measure(EXECUTE, () -> "done"));
+    assertTrue(context.getProfile().getPhases().get("execute").getTimeMillis() >= 0);
   }
 
   @Test
   public void testMeasureExecutesWhenProfilingDisabled() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(false)) {
+    try (UnifiedQueryContext ctx = super.buildContext().build()) {
       assertEquals("done", ctx.measure(EXECUTE, () -> "done"));
       assertNull(ctx.getProfile());
     }
   }
 
   @Test
-  public void testMeasurePropagatesException() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      assertThrows(
-          IOException.class,
-          () ->
-              ctx.measure(
-                  EXECUTE,
-                  () -> {
-                    throw new IOException("test error");
-                  }));
-    }
+  public void testMeasurePropagatesException() {
+    assertThrows(
+        IOException.class,
+        () ->
+            context.measure(
+                EXECUTE,
+                () -> {
+                  throw new IOException("test error");
+                }));
   }
 
   @Test
   public void testFullPipelineProfiling() throws Exception {
-    try (UnifiedQueryContext ctx = buildContext(true)) {
-      RelNode plan = new UnifiedQueryPlanner(ctx).plan("source = opensearch.employees");
-      PreparedStatement stmt = new UnifiedQueryCompiler(ctx).compile(plan);
-      ResultSet rs = ctx.measure(EXECUTE, stmt::executeQuery);
-      assertNotNull(rs);
-
-      QueryProfile profile = ctx.getProfile();
-      assertNotNull(profile);
-      assertTrue(profile.getSummary().getTotalTimeMillis() > 0);
-      assertTrue(profile.getPhases().get("analyze").getTimeMillis() > 0);
-      assertTrue(profile.getPhases().get("optimize").getTimeMillis() > 0);
-      assertTrue(profile.getPhases().get("execute").getTimeMillis() >= 0);
-    }
-  }
-
-  private UnifiedQueryContext buildContext(boolean profiling) {
-    return UnifiedQueryContext.builder()
-        .language(QueryType.PPL)
-        .catalog("opensearch", testSchema)
-        .defaultNamespace("opensearch")
-        .profiling(profiling)
-        .build();
+    RelNode plan = planner.plan("source = catalog.employees");
+    PreparedStatement stmt = new UnifiedQueryCompiler(context).compile(plan);
+    ResultSet rs = context.measure(EXECUTE, stmt::executeQuery);
+
+    QueryProfile profile = context.getProfile();
+    assertTrue(profile.getSummary().getTotalTimeMillis() >= 0);
+    assertTrue(profile.getPhases().get("analyze").getTimeMillis() >= 0);
+    assertTrue(profile.getPhases().get("optimize").getTimeMillis() >= 0);
+    assertTrue(profile.getPhases().get("execute").getTimeMillis() >= 0);
   }
 }

api/src/testFixtures/java/org/opensearch/sql/api/UnifiedQueryTestBase.java

@@ -55,14 +55,20 @@ protected Map<String, Table> getTableMap() {
           }
         };
 
-    context =
-        UnifiedQueryContext.builder()
-            .language(QueryType.PPL)
-            .catalog(DEFAULT_CATALOG, testSchema)
-            .build();
+    context = buildContext().build();
     planner = new UnifiedQueryPlanner(context);
   }
 
+  /**
+   * Creates a pre-configured context builder with test schema. Subclasses can override to customize
+   * context configuration (e.g., enable profiling).
+   */
+  protected UnifiedQueryContext.Builder buildContext() {
+    return UnifiedQueryContext.builder()
+        .language(QueryType.PPL)
+        .catalog(DEFAULT_CATALOG, testSchema);
+  }
+
   @After
   public void tearDown() throws Exception {
     if (context != null) {