Factor JDBC test infrastructure into a fluent API, JdbcAssert.

Signed-off-by: Jacques Nadeau <[email protected]>

sandbox/prototype/sqlparser/src/test/java/org/apache/drill/jdbc/test/JdbcAssert.java

@@ -0,0 +1,121 @@
+/*******************************************************************************
+ * 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 org.apache.drill.jdbc.test;
+
+import com.google.common.base.Function;
+import junit.framework.Assert;
+
+import java.sql.*;
+import java.util.Properties;
+
+/**
+ * Fluent interface for writing JDBC and query-planning tests.
+ */
+public class JdbcAssert {
+  public static One withModel(String model, String schema) {
+    final Properties info = new Properties();
+    info.setProperty("schema", schema);
+    info.setProperty("model", "inline:" + model);
+    return new One(info);
+  }
+
+  static String toString(ResultSet resultSet) throws SQLException {
+    StringBuilder buf = new StringBuilder();
+    while (resultSet.next()) {
+      int n = resultSet.getMetaData().getColumnCount();
+      String sep = "";
+      for (int i = 1; i <= n; i++) {
+        buf.append(sep)
+            .append(resultSet.getMetaData().getColumnLabel(i))
+            .append("=")
+            .append(resultSet.getObject(i));
+        sep = "; ";
+      }
+      buf.append("\n");
+    }
+    return buf.toString();
+  }
+
+  public static class One {
+    private final Properties info;
+    private final ConnectionFactory connectionFactory;
+
+    public One(Properties info) {
+      this.info = info;
+      this.connectionFactory = new ConnectionFactory() {
+        public Connection createConnection() throws Exception {
+          Class.forName("org.apache.drill.jdbc.Driver");
+          return DriverManager.getConnection("jdbc:drill:", One.this.info);
+        }
+      };
+    }
+
+    public Two sql(String sql) {
+      return new Two(connectionFactory, sql);
+    }
+
+    public <T> T withConnection(Function<Connection, T> function)
+        throws Exception {
+      Connection connection = null;
+      try {
+        connection = connectionFactory.createConnection();
+        return function.apply(connection);
+      } finally {
+        if (connection != null) {
+          connection.close();
+        }
+      }
+    }
+  }
+
+  public static class Two {
+    private final ConnectionFactory connectionFactory;
+    private final String sql;
+
+    Two(ConnectionFactory connectionFactory, String sql) {
+      this.connectionFactory = connectionFactory;
+      this.sql = sql;
+    }
+
+    public Two returns(String expected) throws Exception {
+      Connection connection = null;
+      Statement statement = null;
+      try {
+        connection = connectionFactory.createConnection();
+        statement = connection.createStatement();
+        ResultSet resultSet = statement.executeQuery(sql);
+        Assert.assertEquals(expected, JdbcAssert.toString(resultSet));
+        resultSet.close();
+        return this;
+      } finally {
+        if (statement != null) {
+          statement.close();
+        }
+        if (connection != null) {
+          connection.close();
+        }
+      }
+    }
+  }
+
+  private static interface ConnectionFactory {
+    Connection createConnection() throws Exception;
+  }
+}
+
+// End JdbcAssert.java

sandbox/prototype/sqlparser/src/test/java/org/apache/drill/jdbc/test/JdbcTest.java

@@ -17,11 +17,13 @@
  ******************************************************************************/
 package org.apache.drill.jdbc.test;
 
+import com.google.common.base.Function;
+
 import junit.framework.TestCase;
+
 import org.apache.drill.jdbc.DrillTable;
 
 import java.sql.*;
-import java.util.Properties;
 
 /** Unit tests for Drill's JDBC driver. */
 public class JdbcTest extends TestCase {
@@ -58,126 +60,88 @@ public void testLoadDriver() throws ClassNotFoundException {
   public void testConnect() throws Exception {
     Class.forName("org.apache.drill.jdbc.Driver");
     final Connection connection = DriverManager.getConnection(
-        "jdbc:drill:schema=DONUTS;tables=EMP,DEPT");
+        "jdbc:drill:schema=DONUTS");
     connection.close();
   }
 
   /** Load driver, make a connection, prepare a statement. */
-  public void _testPrepare() throws Exception {
-    Class.forName("org.apache.drill.jdbc.Driver");
-    final Connection connection = DriverManager.getConnection(
-        "jdbc:drill:schema=DONUTS;tables=EMP,DEPT");
-    final Statement statement = connection.prepareStatement(
-        "select * from emp where cast(name as varchar(10)) = 'Eric'");
-    statement.close();
-    connection.close();
+  public void testPrepare() throws Exception {
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .withConnection(
+            new Function<Connection, Void>() {
+              public Void apply(Connection connection) {
+                try {
+                  final Statement statement = connection.prepareStatement(
+                      "select * from donuts");
+                  statement.close();
+                  return null;
+                } catch (Exception e) {
+                  throw new RuntimeException(e);
+                }
+              }
+            });
   }
 
   /** Simple query against JSON. */
   public void testSelectJson() throws Exception {
-    Class.forName("org.apache.drill.jdbc.Driver");
-    final Properties info = new Properties();
-    info.setProperty("schema", "DONUTS");
-    info.setProperty("model", "inline:" + MODEL);
-    final Connection connection =
-        DriverManager.getConnection("jdbc:drill:", info);
-    final Statement statement = connection.createStatement();
-    final ResultSet resultSet = statement.executeQuery(
-        "select * from donuts");
-    assertEquals(
-        EXPECTED,
-        toString(resultSet));
-    resultSet.close();
-    statement.close();
-    connection.close();
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select * from donuts")
+        .returns(EXPECTED);
   }
 
   /** Query with project list. No field references yet. */
   public void testProjectConstant() throws Exception {
-    assertSqlReturns(
-        "select 1 + 3 as c from donuts",
-        "C=4\n"
-        + "C=4\n"
-        + "C=4\n"
-        + "C=4\n"
-        + "C=4\n");
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select 1 + 3 as c from donuts")
+        .returns("C=4\n"
+            + "C=4\n"
+            + "C=4\n"
+            + "C=4\n"
+            + "C=4\n");
   }
 
   /** Query that projects an element from the map. */
   public void testProject() throws Exception {
-    assertSqlReturns(
-        "select _MAP['donuts']['ppu'] as ppu from donuts",
-        "PPU=0.55\n"
-        + "PPU=0.69\n"
-        + "PPU=0.55\n"
-        + "PPU=0.69\n"
-        + "PPU=1.0\n");
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select _MAP['donuts']['ppu'] as ppu from donuts")
+        .returns("PPU=0.55\n"
+            + "PPU=0.69\n"
+            + "PPU=0.55\n"
+            + "PPU=0.69\n"
+            + "PPU=1.0\n");
   }
 
   /** Query with subquery, filter, and projection of one real and one
    * nonexistent field from a map field. */
   public void testProjectFilterSubquery() throws Exception {
-    assertSqlReturns(
-        "select d['name'] as name, d['xx'] as xx from (\n"
-        + " select _MAP['donuts'] as d from donuts)\n"
-        + "where cast(d['ppu'] as double) > 0.6",
-        "NAME=Raised; XX=null\n"
-        + "NAME=Filled; XX=null\n"
-        + "NAME=Apple Fritter; XX=null\n");
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select d['name'] as name, d['xx'] as xx from (\n"
+            + " select _MAP['donuts'] as d from donuts)\n"
+            + "where cast(d['ppu'] as double) > 0.6")
+        .returns("NAME=Raised; XX=null\n"
+            + "NAME=Filled; XX=null\n"
+            + "NAME=Apple Fritter; XX=null\n");
   }
 
   /** Query that projects one field. (Disabled; uses sugared syntax.) */
   public void _testProjectNestedFieldSugared() throws Exception {
-    assertSqlReturns(
-        "select donuts.ppu from donuts",
-        "C=4\n"
-        + "C=4\n"
-        + "C=4\n"
-        + "C=4\n"
-        + "C=4\n");
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select donuts.ppu from donuts")
+        .returns("C=4\n"
+            + "C=4\n"
+            + "C=4\n"
+            + "C=4\n"
+            + "C=4\n");
   }
 
   /** Query with filter. No field references yet. */
   public void testFilterConstant() throws Exception {
-    assertSqlReturns(
-        "select * from donuts where 3 > 4",
-        "");
-    assertSqlReturns(
-        "select * from donuts where 3 < 4",
-        EXPECTED);
-  }
-
-  private void assertSqlReturns(String sql, String expected)
-      throws ClassNotFoundException, SQLException {
-    Class.forName("org.apache.drill.jdbc.Driver");
-    final Properties info = new Properties();
-    info.setProperty("schema", "DONUTS");
-    info.setProperty("model", "inline:" + MODEL);
-    final Connection connection =
-        DriverManager.getConnection("jdbc:drill:", info);
-    try (Statement statement = connection.createStatement();
-         ResultSet resultSet = statement.executeQuery(sql)) {
-      assertEquals(expected, toString(resultSet));
-    } finally {
-      connection.close();
-    }
-  }
-
-  static String toString(ResultSet resultSet) throws SQLException {
-    StringBuilder buf = new StringBuilder();
-    while (resultSet.next()) {
-      int n = resultSet.getMetaData().getColumnCount();
-      String sep = "";
-      for (int i = 1; i <= n; i++) {
-        buf.append(sep)
-            .append(resultSet.getMetaData().getColumnLabel(i))
-            .append("=")
-            .append(resultSet.getObject(i));
-        sep = "; ";
-      }
-      buf.append("\n");
-    }
-    return buf.toString();
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select * from donuts where 3 > 4")
+        .returns("");
+    JdbcAssert.withModel(MODEL, "DONUTS")
+        .sql("select * from donuts where 3 < 4")
+        .returns(EXPECTED);
   }
 }