Use TestTable to ensure test tables are cleaned up

trinodb · Jul 9, 2021 · 96be8f2 · 96be8f2
1 parent 6198afb
commit 96be8f2
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 34 deletions.
diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/BaseJdbcConnectorTest.java
@@ -82,7 +82,6 @@
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TOPN_PUSHDOWN;
 import static io.trino.testing.TestingConnectorBehavior.SUPPORTS_TOPN_PUSHDOWN_WITH_VARCHAR;
 import static io.trino.testing.assertions.Assert.assertEventually;
-import static io.trino.testing.sql.TestTable.randomTableSuffix;
 import static java.lang.String.format;
 import static java.util.Locale.ENGLISH;
 import static java.util.concurrent.Executors.newCachedThreadPool;
@@ -1171,17 +1170,16 @@ protected TestView createSleepingView(Duration minimalSleepDuration)
     public void testDeleteWithBigintEqualityPredicate()
     {
         skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE) && hasBehavior(SUPPORTS_ROW_LEVEL_DELETE));
-        String tableName = "test_delete_" + randomTableSuffix();
-        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM region", 5);
-        assertUpdate("DELETE FROM " + tableName + " WHERE regionkey = 1", 1);
-        assertQuery(
-                "SELECT regionkey, name FROM " + tableName,
-                "VALUES "
-                        + "(0, 'AFRICA'),"
-                        + "(2, 'ASIA'),"
-                        + "(3, 'EUROPE'),"
-                        + "(4, 'MIDDLE EAST')");
-        assertUpdate("DROP TABLE " + tableName);
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_delete", "AS SELECT * FROM region")) {
+            assertUpdate("DELETE FROM " + table.getName() + " WHERE regionkey = 1", 1);
+            assertQuery(
+                    "SELECT regionkey, name FROM " + table.getName(),
+                    "VALUES "
+                            + "(0, 'AFRICA'),"
+                            + "(2, 'ASIA'),"
+                            + "(3, 'EUROPE'),"
+                            + "(4, 'MIDDLE EAST')");
+        }
     }
 
     @Test

diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorSmokeTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorSmokeTest.java
@@ -210,17 +210,14 @@ public void testDeleteAllDataFromTable()
     public void testRowLevelDelete()
     {
         skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE) && hasBehavior(SUPPORTS_ROW_LEVEL_DELETE));
-        String tableName = "test_delete_" + randomTableSuffix();
-        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM region", 5);
-
-        assertUpdate("DELETE FROM " + tableName + " WHERE regionkey = 2", 1);
-        assertThat(query("SELECT * FROM " + tableName + " WHERE regionkey = 2"))
-                .returnsEmptyResult();
-        assertThat(query("SELECT cast(regionkey AS integer) FROM " + tableName))
-                .skippingTypesCheck()
-                .matches("VALUES 0, 1, 3, 4");
-
-        assertUpdate("DROP TABLE " + tableName);
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_delete", "AS SELECT * FROM region")) {
+            assertUpdate("DELETE FROM " + table.getName() + " WHERE regionkey = 2", 1);
+            assertThat(query("SELECT * FROM " + table.getName() + " WHERE regionkey = 2"))
+                    .returnsEmptyResult();
+            assertThat(query("SELECT cast(regionkey AS integer) FROM " + table.getName()))
+                    .skippingTypesCheck()
+                    .matches("VALUES 0, 1, 3, 4");
+        }
     }
 
     @Test

diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java
@@ -760,22 +760,20 @@ public void verifySupportsRowLevelDeleteDeclaration()
     public void testDeleteAllDataFromTable()
     {
         skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE) && hasBehavior(SUPPORTS_DELETE));
-        String tableName = "test_delete_" + randomTableSuffix();
-        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM nation", 25);
-        // not using assertUpdate as some connectors provide update count and some not
-        getQueryRunner().execute("DELETE FROM " + tableName);
-        assertQuery("SELECT count(*) FROM " + tableName, "VALUES 0");
-        assertUpdate("DROP TABLE " + tableName);
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_delete", "AS SELECT * FROM region")) {
+            // not using assertUpdate as some connectors provide update count and some not
+            getQueryRunner().execute("DELETE FROM " + table.getName());
+            assertQuery("SELECT count(*) FROM " + table.getName(), "VALUES 0");
+        }
     }
 
     @Test
     public void testRowLevelDelete()
     {
         skipTestUnless(hasBehavior(SUPPORTS_CREATE_TABLE) && hasBehavior(SUPPORTS_ROW_LEVEL_DELETE));
-        String tableName = "test_delete_" + randomTableSuffix();
-        assertUpdate("CREATE TABLE " + tableName + " AS SELECT * FROM nation", 25);
-        assertUpdate("DELETE FROM " + tableName + " WHERE nationkey = 2", 1);
-        assertQuery("SELECT count(*) FROM " + tableName, "VALUES 24");
-        assertUpdate("DROP TABLE " + tableName);
+        try (TestTable table = new TestTable(getQueryRunner()::execute, "test_delete", "AS SELECT * FROM region")) {
+            assertUpdate("DELETE FROM " + table.getName() + " WHERE regionkey = 2", 1);
+            assertQuery("SELECT count(*) FROM " + table.getName(), "VALUES 4");
+        }
     }
 }