From a107ee3ed2c1106fd6e6a9195c3519c86fff0de3 Mon Sep 17 00:00:00 2001
From: Maytas Monsereenusorn <52679095+maytasm@users.noreply.github.com>
Date: Wed, 29 Apr 2020 05:53:25 -1000
Subject: [PATCH] Fix problem when running single integration test using
 -Dit.test= (#9778)

* fix running single it

* fix checksyle
---
 integration-tests/pom.xml                     |  12 ++
 .../org/testng/DruidTestRunnerFactory.java    | 105 ++++++++++++++++++
 2 files changed, 117 insertions(+)
 create mode 100644 integration-tests/src/main/java/org/testng/DruidTestRunnerFactory.java

diff --git a/integration-tests/pom.xml b/integration-tests/pom.xml
index 9c4c429f738..4c2ef70ee52 100644
--- a/integration-tests/pom.xml
+++ b/integration-tests/pom.xml
@@ -405,6 +405,12 @@
                             </execution>
                         </executions>
                         <configuration>
+                            <properties>
+                                <property>
+                                    <name>testrunfactory</name>
+                                    <value>org.testng.DruidTestRunnerFactory</value>
+                                </property>
+                            </properties>
                             <argLine>
                                 -Duser.timezone=UTC
                                 -Dfile.encoding=UTF-8
@@ -457,6 +463,12 @@
                             </execution>
                         </executions>
                         <configuration>
+                            <properties>
+                                <property>
+                                    <name>testrunfactory</name>
+                                    <value>org.testng.DruidTestRunnerFactory</value>
+                                </property>
+                            </properties>
                             <argLine>
                                 -Duser.timezone=UTC
                                 -Dfile.encoding=UTF-8
diff --git a/integration-tests/src/main/java/org/testng/DruidTestRunnerFactory.java b/integration-tests/src/main/java/org/testng/DruidTestRunnerFactory.java
new file mode 100644
index 00000000000..daafff44c75
--- /dev/null
+++ b/integration-tests/src/main/java/org/testng/DruidTestRunnerFactory.java
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package /*CHECKSTYLE.OFF: PackageName*/org.testng/*CHECKSTYLE.ON: PackageName*/;
+
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.testing.utils.SuiteListener;
+import org.testng.internal.IConfiguration;
+import org.testng.internal.Systematiser;
+import org.testng.internal.annotations.IAnnotationFinder;
+import org.testng.xml.XmlTest;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * This class must be in package org.testng to access protected methods like TestNG.getDefault().getConfiguration()
+ */
+public class DruidTestRunnerFactory implements ITestRunnerFactory
+{
+  private static final Logger LOG = new Logger(DruidTestRunnerFactory.class);
+  private static final SuiteListener SUITE_LISTENER = new SuiteListener();
+
+  @Override
+  public TestRunner newTestRunner(ISuite suite,
+                                  XmlTest test,
+                                  Collection<IInvokedMethodListener> methodListeners,
+                                  List<IClassListener> classListeners)
+  {
+    IConfiguration configuration = TestNG.getDefault().getConfiguration();
+    String outputDirectory = suite.getOutputDirectory();
+    IAnnotationFinder annotationFinder = configuration.getAnnotationFinder();
+    Boolean skipFailedInvocationCounts = suite.getXmlSuite().skipFailedInvocationCounts();
+    return new DruidTestRunner(
+        configuration,
+        suite,
+        test,
+        outputDirectory,
+        annotationFinder,
+        skipFailedInvocationCounts,
+        methodListeners,
+        classListeners
+    );
+  }
+
+  private static class DruidTestRunner extends TestRunner
+  {
+    DruidTestRunner(
+        IConfiguration configuration,
+        ISuite suite,
+        XmlTest test,
+        String outputDirectory,
+        IAnnotationFinder finder,
+        boolean skipFailedInvocationCounts,
+        Collection<IInvokedMethodListener> methodListeners,
+        List<IClassListener> classListeners
+    )
+    {
+      super(configuration, suite, test, outputDirectory, finder, skipFailedInvocationCounts, methodListeners, classListeners, Systematiser.getComparator(), Collections.emptyMap());
+    }
+
+    @Override
+    public void run()
+    {
+      try {
+        // IntegrationTestSuite is run when -Dit.test is not specify in maven command.
+        // IntegrationTestSuite uses the configuration from integration-tests/src/test/resources/testng.xml
+        // which already handle suite onStart and onFinish automatically
+        if (!"IntegrationTestSuite".equals(getSuite().getName())) {
+          SUITE_LISTENER.onStart(getSuite());
+        }
+        runTests();
+      }
+      finally {
+        // IntegrationTestSuite uses the configuration from integration-tests/src/test/resources/testng.xml
+        // which already handle suite onStart and onFinish automatically
+        if (!"IntegrationTestSuite".equals(getSuite().getName())) {
+          SUITE_LISTENER.onFinish(getSuite());
+        }
+      }
+    }
+
+    private void runTests()
+    {
+      super.run();
+    }
+  }
+}