Fix UTs and integration tests

apache · Sep 3, 2023 · 7d783ca · 7d783ca
1 parent 2808f64
commit 7d783ca
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 5 deletions.
diff --git a/pinot-broker/pom.xml b/pinot-broker/pom.xml
@@ -152,6 +152,12 @@
       <artifactId>testng</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-query-event-listener</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>

diff --git a/...er/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java b/...er/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java
@@ -32,6 +32,7 @@
 import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
 import org.apache.pinot.spi.env.PinotConfiguration;
 import org.apache.pinot.spi.metrics.PinotMetricUtils;
+import org.apache.pinot.spi.queryeventlistener.PinotBrokerQueryEventListenerUtils;
 import org.apache.pinot.spi.trace.RequestContext;
 import org.apache.pinot.spi.trace.Tracing;
 import org.apache.pinot.spi.utils.BytesUtils;
@@ -183,7 +184,8 @@ public void testBrokerRequestHandler()
     SingleConnectionBrokerRequestHandler requestHandler =
         new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY,
             null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()),
-            null, null, mock(ServerRoutingStatsManager.class), null);
+            null, null, mock(ServerRoutingStatsManager.class),
+                PinotBrokerQueryEventListenerUtils.getBrokerQueryEventListener());
 
     long randNum = RANDOM.nextLong();
     byte[] randBytes = new byte[12];
@@ -211,7 +213,8 @@ public void testBrokerRequestHandlerWithAsFunction()
     SingleConnectionBrokerRequestHandler requestHandler =
         new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY,
             null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()),
-            null, null, mock(ServerRoutingStatsManager.class), null);
+            null, null, mock(ServerRoutingStatsManager.class),
+                PinotBrokerQueryEventListenerUtils.getBrokerQueryEventListener());
     long currentTsMin = System.currentTimeMillis();
     JsonNode request = JsonUtils.stringToJsonNode(
         "{\"sql\":\"SELECT now() as currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') as firstDayOf2020\"}");
@@ -418,7 +421,8 @@ public void testExplainPlanLiteralOnly()
     SingleConnectionBrokerRequestHandler requestHandler =
         new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY,
             null, null, new BrokerMetrics("", PinotMetricUtils.getPinotMetricsRegistry(), true, Collections.emptySet()),
-            null, null, mock(ServerRoutingStatsManager.class), null);
+            null, null, mock(ServerRoutingStatsManager.class),
+                PinotBrokerQueryEventListenerUtils.getBrokerQueryEventListener());
 
     // Test 1: select constant
     JsonNode request = JsonUtils.stringToJsonNode("{\"sql\":\"EXPLAIN PLAN FOR SELECT 1.5, 'test'\"}");

diff --git a/.../test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java b/.../test/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandlerTest.java
@@ -29,6 +29,7 @@
 import org.apache.pinot.common.config.provider.TableCache;
 import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.queryeventlistener.PinotBrokerQueryEventListenerUtils;
 import org.apache.pinot.spi.trace.DefaultRequestContext;
 import org.apache.pinot.spi.trace.RequestContext;
 import org.apache.pinot.spi.utils.CommonConstants;
@@ -65,8 +66,9 @@ public void setUp() {
     _config.setProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, "12345");
     _accessControlFactory = new AllowAllAccessControlFactory();
     _requestHandler =
-        new MultiStageBrokerRequestHandler(_config, "Broker_localhost", _routingManager, _accessControlFactory,
-            _queryQuotaManager, _tableCache, _brokerMetrics, null);
+        new MultiStageBrokerRequestHandler(_config, "Broker_localhost", _routingManager,
+                _accessControlFactory, _queryQuotaManager, _tableCache, _brokerMetrics,
+                PinotBrokerQueryEventListenerUtils.getBrokerQueryEventListener());
   }
 
   @Test

diff --git a/pinot-integration-test-base/pom.xml b/pinot-integration-test-base/pom.xml
@@ -168,6 +168,11 @@
       <artifactId>pinot-yammer</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-query-event-listener</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>org.testng</groupId>
       <artifactId>testng</artifactId>

diff --git a/pinot-integration-tests/pom.xml b/pinot-integration-tests/pom.xml
@@ -339,6 +339,11 @@
       <version>${project.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-query-event-listener</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>org.testcontainers</groupId>
       <artifactId>testcontainers</artifactId>

diff --git a/pinot-tools/pom.xml b/pinot-tools/pom.xml
@@ -230,6 +230,11 @@
       <version>${project.version}</version>
       <scope>runtime</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-query-event-listener</artifactId>
+      <version>${project.version}</version>
+    </dependency>
     <dependency>
       <groupId>commons-cli</groupId>
       <artifactId>commons-cli</artifactId>