Add helper to simplify tests & add canonical string

Author: aharpervc

src/test_utils.rs

@@ -151,6 +151,8 @@ impl TestedDialects {
     ///
     /// 2. re-serializing the result of parsing `sql` produces the same
     ///    `canonical` sql string
+    ///
+    /// For multiple statements, use [`multiple_statements_parse_to`].
     pub fn one_statement_parses_to(&self, sql: &str, canonical: &str) -> Statement {
         let mut statements = self.parse_sql_statements(sql).expect(sql);
         assert_eq!(statements.len(), 1);
@@ -166,6 +168,32 @@ impl TestedDialects {
         only_statement
     }
 
+    /// The same as [`one_statement_parses_to`] but it works for a multiple statements
+    pub fn multiple_statements_parse_to(
+        &self,
+        sql: &str,
+        statement_count: usize,
+        canonical: &str,
+    ) -> Vec<Statement> {
+        let statements = self.parse_sql_statements(sql).expect(sql);
+        assert_eq!(statements.len(), statement_count);
+
+        if !canonical.is_empty() && sql != canonical {
+            assert_eq!(self.parse_sql_statements(canonical).unwrap(), statements);
+        } else {
+            assert_eq!(
+                sql,
+                statements
+                    .iter()
+                    .map(|s| s.to_string())
+                    .collect::<Vec<_>>()
+                    .join("; ")
+            );
+        }
+
+        statements
+    }
+
     /// Ensures that `sql` parses as an [`Expr`], and that
     /// re-serializing the parse result produces canonical
     pub fn expr_parses_to(&self, sql: &str, canonical: &str) -> Expr {

tests/sqlparser_mssql.rs

@@ -2161,28 +2161,27 @@ fn parse_print() {
 #[test]
 fn parse_mssql_go_keyword() {
     let single_go_keyword = "USE some_database;\nGO";
-    let stmts = ms().parse_sql_statements(single_go_keyword).unwrap();
-    assert_eq!(stmts.len(), 2);
-    assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }),);
+    let stmts = ms().multiple_statements_parse_to(single_go_keyword, 2, "USE some_database\nGO");
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
 
     let go_with_count = "SELECT 1;\nGO 5";
-    let stmts = ms().parse_sql_statements(go_with_count).unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(go_with_count, 2, "SELECT 1\nGO 5");
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(5) }));
 
     let go_statement_delimiter = "SELECT 1\nGO";
-    let stmts = ms().parse_sql_statements(go_statement_delimiter).unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(go_statement_delimiter, 2, "SELECT 1; \nGO");
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
 
     let bare_go = "GO";
-    let stmts = ms().parse_sql_statements(bare_go).unwrap();
-    assert_eq!(stmts.len(), 1);
-    assert_eq!(stmts[0], Statement::Go(GoStatement { count: None }));
+    let stmt = ms().one_statement_parses_to(bare_go, "GO");
+    assert_eq!(stmt, Statement::Go(GoStatement { count: None }));
 
     let go_then_statements = "/* whitespace */ GO\nRAISERROR('This is a test', 16, 1);";
-    let stmts = ms().parse_sql_statements(go_then_statements).unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(
+        go_then_statements,
+        2,
+        "GO\nRAISERROR('This is a test', 16, 1)",
+    );
     assert_eq!(stmts[0], Statement::Go(GoStatement { count: None }));
     assert_eq!(
         stmts[1],
@@ -2198,41 +2197,45 @@ fn parse_mssql_go_keyword() {
     );
 
     let multiple_gos = "SELECT 1;\nGO 5\nSELECT 2;\n  GO";
-    let stmts = ms().parse_sql_statements(multiple_gos).unwrap();
-    assert_eq!(stmts.len(), 4);
+    let stmts = ms().multiple_statements_parse_to(multiple_gos, 4, "SELECT 1\nGO 5\nSELECT 2\nGO");
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(5) }));
     assert_eq!(stmts[3], Statement::Go(GoStatement { count: None }));
 
     let single_line_comment_preceding_go = "USE some_database; -- okay\nGO";
-    let stmts = ms()
-        .parse_sql_statements(single_line_comment_preceding_go)
-        .unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(
+        single_line_comment_preceding_go,
+        2,
+        "USE some_database\nGO",
+    );
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
 
     let multi_line_comment_preceding_go = "USE some_database; /* okay */\nGO";
-    let stmts = ms()
-        .parse_sql_statements(multi_line_comment_preceding_go)
-        .unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(
+        multi_line_comment_preceding_go,
+        2,
+        "USE some_database\nGO",
+    );
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
 
     let single_line_comment_following_go = "USE some_database;\nGO -- okay";
-    let stmts = ms().parse_sql_statements(single_line_comment_following_go).unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(
+        single_line_comment_following_go,
+        2,
+        "USE some_database\nGO",
+    );
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
 
     let multi_line_comment_following = "USE some_database;\nGO/* okay */42";
-    let stmts = ms()
-        .parse_sql_statements(multi_line_comment_following)
-        .unwrap();
-    assert_eq!(stmts.len(), 2);
+    let stmts = ms().multiple_statements_parse_to(
+        multi_line_comment_following,
+        2,
+        "USE some_database\nGO 42",
+    );
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(42) }));
 
     let actually_column_alias = "SELECT NULL GO";
-    let stmts = ms().parse_sql_statements(actually_column_alias).unwrap();
-    assert_eq!(stmts.len(), 1);
-    match &stmts[0] {
+    let stmt = ms().one_statement_parses_to(actually_column_alias, "SELECT NULL AS GO");
+    match &stmt {
         Statement::Query(query) => {
             let select = query.body.as_select().unwrap();
             assert_eq!(