From cc781512144ad395d48e1dd3e86bab9bd692aa38 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Grzegorz=20Kokosi=C5=84ski?= <grzegorz@starburstdata.com>
Date: Thu, 20 Aug 2020 14:10:37 +0200
Subject: [PATCH] Use generic trueFalse data provider for timestamp tests

---
 .../java/io/prestosql/testing/TestngUtils.java     |  8 ++++++++
 .../postgresql/TestPostgreSqlTypeMapping.java      | 14 +++-----------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/presto-main/src/main/java/io/prestosql/testing/TestngUtils.java b/presto-main/src/main/java/io/prestosql/testing/TestngUtils.java
index ab648df9e3b6..d14db4f1dd54 100644
--- a/presto-main/src/main/java/io/prestosql/testing/TestngUtils.java
+++ b/presto-main/src/main/java/io/prestosql/testing/TestngUtils.java
@@ -13,6 +13,8 @@
  */
 package io.prestosql.testing;
 
+import org.testng.annotations.DataProvider;
+
 import java.util.ArrayList;
 import java.util.stream.Collector;
 
@@ -31,4 +33,10 @@ private TestngUtils() {}
                 },
                 builder -> builder.toArray(new Object[][] {}));
     }
+
+    @DataProvider
+    public static Object[][] trueFalse()
+    {
+        return new Object[][] {{true}, {false}};
+    }
 }
diff --git a/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java b/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java
index cf4efb16b33f..dce5ff368ed1 100644
--- a/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java
+++ b/presto-postgresql/src/test/java/io/prestosql/plugin/postgresql/TestPostgreSqlTypeMapping.java
@@ -26,6 +26,7 @@
 import io.prestosql.testing.AbstractTestQueryFramework;
 import io.prestosql.testing.QueryRunner;
 import io.prestosql.testing.TestingSession;
+import io.prestosql.testing.TestngUtils;
 import io.prestosql.testing.datatype.CreateAndInsertDataSetup;
 import io.prestosql.testing.datatype.CreateAndPrestoInsertDataSetup;
 import io.prestosql.testing.datatype.CreateAsSelectDataSetup;
@@ -1086,7 +1087,7 @@ public Object[][] testTimestampDataProvider()
         };
     }
 
-    @Test(dataProvider = "testTimestampWithTimeZoneDataProvider")
+    @Test(dataProvider = "trueFalse", dataProviderClass = TestngUtils.class)
     public void testTimestampWithTimeZone(boolean insertWithPresto)
     {
         DataSetup dataSetup;
@@ -1152,7 +1153,7 @@ public void testCreateTableWithInvalidTimestampWithTimeZone()
                 "Unsupported column type: timestamp\\(7\\) with time zone");
     }
 
-    @Test(dataProvider = "testTimestampWithTimeZoneDataProvider")
+    @Test(dataProvider = "trueFalse", dataProviderClass = TestngUtils.class)
     public void testArrayTimestampWithTimeZone(boolean insertWithPresto)
     {
         DataSetup dataSetup;
@@ -1198,15 +1199,6 @@ public void testArrayTimestampWithTimeZone(boolean insertWithPresto)
         tests.execute(getQueryRunner(), sessionWithArrayAsArray(), dataSetup);
     }
 
-    @DataProvider
-    public Object[][] testTimestampWithTimeZoneDataProvider()
-    {
-        return new Object[][] {
-                {true},
-                {false},
-        };
-    }
-
     @Test
     public void testJson()
     {