intellij-language-annotation: add annotation on tests (apache#13793)

shounakmk219 · Aug 18, 2024 · 1861158 · 1861158
1 parent cc14cbe
commit 1861158
Show file tree

Hide file tree

Showing 7 changed files with 61 additions and 38 deletions.
diff --git a/pinot-core/pom.xml b/pinot-core/pom.xml
@@ -151,6 +151,11 @@
       <artifactId>pinot-yammer</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.jetbrains</groupId>
+      <artifactId>annotations</artifactId>
+      <scope>test</scope>
+    </dependency>
 
     <!-- Lucene dependencies -->
     <dependency>

diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java
@@ -63,6 +63,7 @@
 import org.apache.pinot.spi.utils.ReadMode;
 import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.intellij.lang.annotations.Language;
 
 import static org.mockito.Mockito.mock;
 
@@ -76,6 +77,7 @@ public abstract class BaseQueriesTest {
   protected static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(2);
   protected static final BrokerMetrics BROKER_METRICS = mock(BrokerMetrics.class);
 
+  @Language(value = "sql", prefix = "select * from table")
   protected abstract String getFilter();
 
   protected abstract IndexSegment getIndexSegment();
@@ -91,7 +93,7 @@ protected List<List<IndexSegment>> getDistinctInstances() {
    * <p>Use this to test a single operator.
    */
   @SuppressWarnings({"rawtypes", "unchecked"})
-  protected <T extends Operator> T getOperator(String query) {
+  protected <T extends Operator> T getOperator(@Language("sql") String query) {
     PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
     PinotQuery serverPinotQuery = GapfillUtils.stripGapfill(pinotQuery);
     QueryContext queryContext = QueryContextConverterUtils.getQueryContext(serverPinotQuery);
@@ -103,7 +105,7 @@ protected <T extends Operator> T getOperator(String query) {
    * <p>Use this to test a single operator.
    */
   @SuppressWarnings("rawtypes")
-  protected <T extends Operator> T getOperatorWithFilter(String query) {
+  protected <T extends Operator> T getOperatorWithFilter(@Language("sql") String query) {
     return getOperator(query + getFilter());
   }
 
@@ -117,7 +119,7 @@ protected <T extends Operator> T getOperatorWithFilter(String query) {
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  protected BrokerResponseNative getBrokerResponse(String query) {
+  protected BrokerResponseNative getBrokerResponse(@Language("sql") String query) {
     return getBrokerResponse(query, PLAN_MAKER);
   }
 
@@ -131,7 +133,7 @@ protected BrokerResponseNative getBrokerResponse(String query) {
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  protected BrokerResponseNative getBrokerResponseWithFilter(String query) {
+  protected BrokerResponseNative getBrokerResponseWithFilter(@Language("sql") String query) {
     return getBrokerResponse(query + getFilter());
   }
 
@@ -145,7 +147,7 @@ protected BrokerResponseNative getBrokerResponseWithFilter(String query) {
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  protected BrokerResponseNative getBrokerResponse(String query, PlanMaker planMaker) {
+  protected BrokerResponseNative getBrokerResponse(@Language("sql") String query, PlanMaker planMaker) {
     return getBrokerResponse(query, planMaker, null);
   }
 
@@ -159,7 +161,8 @@ protected BrokerResponseNative getBrokerResponse(String query, PlanMaker planMak
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  protected BrokerResponseNative getBrokerResponse(String query, @Nullable Map<String, String> extraQueryOptions) {
+  protected BrokerResponseNative getBrokerResponse(
+      @Language("sql") String query, @Nullable Map<String, String> extraQueryOptions) {
     return getBrokerResponse(query, PLAN_MAKER, extraQueryOptions);
   }
 
@@ -173,7 +176,7 @@ protected BrokerResponseNative getBrokerResponse(String query, @Nullable Map<Str
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  private BrokerResponseNative getBrokerResponse(String query, PlanMaker planMaker,
+  private BrokerResponseNative getBrokerResponse(@Language("sql") String query, PlanMaker planMaker,
       @Nullable Map<String, String> extraQueryOptions) {
     PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
     if (extraQueryOptions != null) {
@@ -265,8 +268,8 @@ protected BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, Br
    * This can be particularly useful to test statistical aggregation functions.
    * @see StatisticalQueriesTest for an example use case.
    */
-  protected BrokerResponseNative getBrokerResponseForOptimizedQuery(String query, @Nullable TableConfig config,
-      @Nullable Schema schema) {
+  protected BrokerResponseNative getBrokerResponseForOptimizedQuery(
+      @Language("sql") String query, @Nullable TableConfig config, @Nullable Schema schema) {
     PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
     OPTIMIZER.optimize(pinotQuery, config, schema);
     return getBrokerResponse(pinotQuery, PLAN_MAKER);

diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/FluentQueryTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/FluentQueryTest.java
@@ -52,6 +52,7 @@
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.FileFormat;
 import org.apache.pinot.spi.utils.ReadMode;
+import org.intellij.lang.annotations.Language;
 import org.testng.Assert;
 
 
@@ -192,7 +193,7 @@ protected void processSegments() {
       _segmentContents.clear();
     }
 
-    public QueryExecuted whenQuery(String query) {
+    public QueryExecuted whenQuery(@Language("sql") String query) {
       processSegments();
       BrokerResponseNative brokerResponse = _baseQueriesTest.getBrokerResponse(query, _extraQueryOptions);
       return new QueryExecuted(_baseQueriesTest, brokerResponse, _extraQueryOptions);
@@ -319,7 +320,7 @@ public QueryExecuted withNullHandling(boolean enabled) {
       return this;
     }
 
-    public QueryExecuted whenQuery(String query) {
+    public QueryExecuted whenQuery(@Language("sql") String query) {
       BrokerResponseNative brokerResponse = _baseQueriesTest.getBrokerResponse(query);
       return new QueryExecuted(_baseQueriesTest, brokerResponse, _extraQueryOptions);
     }

diff --git a/...est-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java b/...est-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java
@@ -67,6 +67,7 @@
 import org.apache.pinot.spi.utils.builder.TableNameBuilder;
 import org.apache.pinot.tools.utils.KafkaStarterUtils;
 import org.apache.pinot.util.TestUtils;
+import org.intellij.lang.annotations.Language;
 import org.testng.Assert;
 
 
@@ -723,15 +724,15 @@ protected void resetTable(String tableName, TableType tableType, @Nullable Strin
   /**
    * Run equivalent Pinot and H2 query and compare the results.
    */
-  protected void testQuery(String query)
+  protected void testQuery(@Language("sql") String query)
       throws Exception {
     testQuery(query, query);
   }
 
   /**
    * Run equivalent Pinot and H2 query and compare the results.
    */
-  protected void testQuery(String pinotQuery, String h2Query)
+  protected void testQuery(@Language("sql") String pinotQuery, String h2Query)
       throws Exception {
     ClusterIntegrationTestUtils.testQuery(pinotQuery, getBrokerBaseApiUrl(), getPinotConnection(), h2Query,
         getH2Connection(), null, getExtraQueryProperties(), useMultiStageQueryEngine());
@@ -740,7 +741,7 @@ protected void testQuery(String pinotQuery, String h2Query)
   /**
    * Run equivalent Pinot and H2 query and compare the results.
    */
-  protected void testQueryWithMatchingRowCount(String pinotQuery, String h2Query)
+  protected void testQueryWithMatchingRowCount(@Language("sql") String pinotQuery, @Language("sql") String h2Query)
       throws Exception {
     ClusterIntegrationTestUtils.testQueryWithMatchingRowCount(pinotQuery, getBrokerBaseApiUrl(), getPinotConnection(),
         h2Query, getH2Connection(), null, getExtraQueryProperties(), useMultiStageQueryEngine());

diff --git a/...st-base/src/test/java/org/apache/pinot/integration/tests/ClusterIntegrationTestUtils.java b/...st-base/src/test/java/org/apache/pinot/integration/tests/ClusterIntegrationTestUtils.java
@@ -83,6 +83,7 @@
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
 import org.apache.pinot.tools.utils.ExplainPlanUtils;
 import org.apache.pinot.tools.utils.KafkaStarterUtils;
+import org.intellij.lang.annotations.Language;
 import org.testng.Assert;
 
 
@@ -657,17 +658,18 @@ private static Object generateRandomValue(Schema.Type fieldType) {
   /**
    * Run equivalent Pinot and H2 query and compare the results.
    */
-  static void testQuery(String pinotQuery, String queryResourceUrl, org.apache.pinot.client.Connection pinotConnection,
-      String h2Query, Connection h2Connection)
+  static void testQuery(@Language("sql") String pinotQuery, String queryResourceUrl,
+      org.apache.pinot.client.Connection pinotConnection, @Language("sql") String h2Query, Connection h2Connection)
       throws Exception {
     testQuery(pinotQuery, queryResourceUrl, pinotConnection, h2Query, h2Connection, null);
   }
 
   /**
    * Run equivalent Pinot and H2 query and compare the results.
    */
-  static void testQuery(String pinotQuery, String queryResourceUrl, org.apache.pinot.client.Connection pinotConnection,
-      String h2Query, Connection h2Connection, @Nullable Map<String, String> headers)
+  static void testQuery(@Language("sql") String pinotQuery, String queryResourceUrl,
+      org.apache.pinot.client.Connection pinotConnection, @Language("sql") String h2Query, Connection h2Connection,
+      @Nullable Map<String, String> headers)
       throws Exception {
     testQuery(pinotQuery, queryResourceUrl, pinotConnection, h2Query, h2Connection, headers, null, false);
   }
@@ -676,8 +678,8 @@ static void testQuery(String pinotQuery, String queryResourceUrl, org.apache.pin
    * Compare # of rows in pinot and H2 only. Succeed if # of rows matches. Note this only applies to non-aggregation
    * query.
    */
-  static void testQueryWithMatchingRowCount(String pinotQuery, String queryResourceUrl,
-      org.apache.pinot.client.Connection pinotConnection, String h2Query, Connection h2Connection,
+  static void testQueryWithMatchingRowCount(@Language("sql") String pinotQuery, String queryResourceUrl,
+      org.apache.pinot.client.Connection pinotConnection, @Language("sql") String h2Query, Connection h2Connection,
       @Nullable Map<String, String> headers, @Nullable Map<String, String> extraJsonProperties,
       boolean useMultiStageQueryEngine)
       throws Exception {
@@ -689,9 +691,10 @@ static void testQueryWithMatchingRowCount(String pinotQuery, String queryResourc
     }
   }
 
-  static void testQuery(String pinotQuery, String queryResourceUrl, org.apache.pinot.client.Connection pinotConnection,
-      String h2Query, Connection h2Connection, @Nullable Map<String, String> headers,
-      @Nullable Map<String, String> extraJsonProperties, boolean useMultiStageQueryEngine) {
+  static void testQuery(@Language("sql") String pinotQuery, String queryResourceUrl,
+      org.apache.pinot.client.Connection pinotConnection, @Language("sql") String h2Query, Connection h2Connection,
+      @Nullable Map<String, String> headers, @Nullable Map<String, String> extraJsonProperties,
+      boolean useMultiStageQueryEngine) {
     try {
       testQueryInternal(pinotQuery, queryResourceUrl, pinotConnection, h2Query, h2Connection, headers,
           extraJsonProperties, useMultiStageQueryEngine, false, false);
@@ -700,8 +703,8 @@ static void testQuery(String pinotQuery, String queryResourceUrl, org.apache.pin
     }
   }
 
-  static void testQueryViaController(String pinotQuery, String queryResourceUrl,
-      org.apache.pinot.client.Connection pinotConnection, String h2Query, Connection h2Connection,
+  static void testQueryViaController(@Language("sql") String pinotQuery, String queryResourceUrl,
+      org.apache.pinot.client.Connection pinotConnection, @Language("sql") String h2Query, Connection h2Connection,
       @Nullable Map<String, String> headers, @Nullable Map<String, String> extraJsonProperties,
       boolean useMultiStageQueryEngine) {
     try {
@@ -712,7 +715,7 @@ static void testQueryViaController(String pinotQuery, String queryResourceUrl,
     }
   }
 
-  private static void testQueryInternal(String pinotQuery, String queryResourceUrl,
+  private static void testQueryInternal(@Language("sql") String pinotQuery, String queryResourceUrl,
       org.apache.pinot.client.Connection pinotConnection, String h2Query, Connection h2Connection,
       @Nullable Map<String, String> headers, @Nullable Map<String, String> extraJsonProperties,
       boolean useMultiStageQueryEngine, boolean matchingRowCount, boolean viaController)
@@ -843,8 +846,8 @@ private static void testQueryInternal(String pinotQuery, String queryResourceUrl
     }
   }
 
-  private static String getExplainPlan(String pinotQuery, String brokerUrl, @Nullable Map<String, String> headers,
-      @Nullable Map<String, String> extraJsonProperties)
+  private static String getExplainPlan(@Language("sql") String pinotQuery, String brokerUrl,
+      @Nullable Map<String, String> headers, @Nullable Map<String, String> extraJsonProperties)
       throws Exception {
     JsonNode explainPlanForResponse =
         ClusterTest.postQuery("explain plan for " + pinotQuery, getBrokerQueryApiUrl(brokerUrl, false), headers,
@@ -917,7 +920,7 @@ private static int getH2ExpectedValues(Set<String> expectedValues, List<String>
 
   private static void comparePinotResultsWithExpectedValues(Set<String> expectedValues,
       List<String> expectedOrderByValues, org.apache.pinot.client.ResultSet connectionResultSet,
-      Set<String> orderByColumns, String pinotQuery, int h2NumRows, long pinotNumRecordsSelected) {
+      Set<String> orderByColumns, @Language("sql") String pinotQuery, int h2NumRows, long pinotNumRecordsSelected) {
 
     int pinotNumRows = connectionResultSet.getRowCount();
     // No record selected in H2
@@ -1025,7 +1028,7 @@ public static boolean isParsableDouble(String input) {
     }
   }
 
-  public static boolean fuzzyCompare(String h2Value, String brokerValue, String connectionValue) {
+  public static boolean fuzzyCompare(@Language("sql") String h2Value, String brokerValue, String connectionValue) {
     if (("null".equals(h2Value) || h2Value == null)
         && ("null".equals(brokerValue) || brokerValue == null)
         && ("null".equals(connectionValue) || connectionValue == null)) {
@@ -1055,7 +1058,8 @@ public static boolean fuzzyCompare(String h2Value, String brokerValue, String co
     return error;
   }
 
-  private static void failure(String pinotQuery, String h2Query, @Nullable Exception e) {
+  private static void failure(@Language("sql") String pinotQuery, @Language("sql") String h2Query,
+      @Nullable Exception e) {
     String failureMessage = "Caught exception while testing query!";
     failure(pinotQuery, h2Query, failureMessage, e);
   }
@@ -1068,7 +1072,8 @@ private static void failure(String pinotQuery, String h2Query, @Nullable Excepti
    * @param failureMessage Failure message
    * @param e Exception
    */
-  private static void failure(String pinotQuery, String h2Query, String failureMessage, @Nullable Exception e) {
+  private static void failure(@Language("sql") String pinotQuery, @Language("sql") String h2Query,
+      String failureMessage, @Nullable Exception e) {
     failureMessage += "\nPinot query: " + pinotQuery + "\nH2 query: " + h2Query;
     if (e == null) {
       Assert.fail(failureMessage);

diff --git a/...ntegration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryGenerator.java b/...ntegration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryGenerator.java
@@ -40,6 +40,7 @@
 import org.apache.avro.generic.GenericRecord;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.pinot.spi.utils.JsonUtils;
+import org.intellij.lang.annotations.Language;
 
 
 /**
@@ -376,16 +377,16 @@ private PredicateQueryFragment generatePredicate() {
 
   public interface Query {
 
-    String generatePinotQuery();
+    @Language("sql") String generatePinotQuery();
 
-    String generateH2Query();
+    @Language("sql") String generateH2Query();
   }
 
   private interface QueryFragment {
 
-    String generatePinotQuery();
+    @Language("sql") String generatePinotQuery();
 
-    String generateH2Query();
+    @Language("sql") String generateH2Query();
   }
 
   /**
@@ -574,21 +575,23 @@ public String generateH2Query() {
    * Most basic query fragment.
    */
   private static class StringQueryFragment implements QueryFragment {
+    @Language("sql")
     final String _pinotQuery;
+    @Language("sql")
     final String _h2Query;
 
     /**
      * Constructor with same Pinot and H2 query fragment.
      */
-    StringQueryFragment(String query) {
+    StringQueryFragment(@Language("sql") String query) {
       _pinotQuery = query;
       _h2Query = query;
     }
 
     /**
      * Constructor for <code>StringQueryFragment</code> with different Pinot and H2 query fragment.
      */
-    StringQueryFragment(String pinotQuery, String h2Query) {
+    StringQueryFragment(@Language("sql") String pinotQuery, @Language("sql") String h2Query) {
       _pinotQuery = pinotQuery;
       _h2Query = h2Query;
     }

diff --git a/pom.xml b/pom.xml
@@ -1606,6 +1606,11 @@
         <artifactId>surefire-testng</artifactId>
         <version>${surefire.version}</version>
       </dependency>
+      <dependency>
+        <groupId>org.jetbrains</groupId>
+        <artifactId>annotations</artifactId>
+        <version>24.0.1</version>
+      </dependency>
       <dependency>
         <groupId>com.h2database</groupId>
         <artifactId>h2</artifactId>