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,20 @@ 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

@@ -2075,28 +2075,23 @@ 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],
@@ -2112,41 +2107,29 @@ 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!(