[parser] parser: optimize restore unit test (pingcap#75)

xhebox · Oct 8, 2021 · 1dac708 · 1dac708
1 parent 652aeab
commit 1dac708
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 53 deletions.
diff --git a/parser/ast/expressions_test.go b/parser/ast/expressions_test.go
@@ -14,10 +14,7 @@
 package ast_test
 
 import (
-	"strings"
-
 	. "github.com/pingcap/check"
-	"github.com/pingcap/parser"
 	. "github.com/pingcap/parser/ast"
 	_ "github.com/pingcap/tidb/types/parser_driver"
 )
@@ -105,61 +102,42 @@ func (tc *testExpressionsSuite) TestExpresionsVisitorCover(c *C) {
 	}
 }
 
-type exprTestCase struct {
-	sourceSQL string
-	expectSQL string
-}
-
-func (tc *testExpressionsSuite) createTestCase4UnaryOperationExpr() []exprTestCase {
-	return []exprTestCase{
-		{"select ++1", "SELECT ++1"},
-		{"select --1", "SELECT --1"},
-		{"select -+1", "SELECT -+1"},
-		{"select -1", "SELECT -1"},
+func (tc *testExpressionsSuite) TestUnaryOperationExprRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"++1", "++1"},
+		{"--1", "--1"},
+		{"-+1", "-+1"},
+		{"-1", "-1"},
 	}
-}
-
-func (tc *testExpressionsSuite) createTestCase4ColumnNameExpr() []exprTestCase {
-	return []exprTestCase{
-		{"select abc", "SELECT `abc`"},
-		{"select `abc`", "SELECT `abc`"},
-		{"select `ab``c`", "SELECT `ab``c`"},
-		{"select sabc.tABC", "SELECT `sabc`.`tABC`"},
-		{"select dabc.sabc.tabc", "SELECT `dabc`.`sabc`.`tabc`"},
-		{"select dabc.`sabc`.tabc", "SELECT `dabc`.`sabc`.`tabc`"},
-		{"select `dABC`.`sabc`.tabc", "SELECT `dABC`.`sabc`.`tabc`"},
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr
 	}
+	RunNodeRestoreTest(c, testCases, "select %s", extractNodeFunc)
 }
 
-func (tc *testExpressionsSuite) createTestCase4IsNullExpr() []exprTestCase {
-	return []exprTestCase{
-		{"select a is null", "SELECT `a` IS NULL"},
-		{"select a is not null", "SELECT `a` IS NOT NULL"},
+func (tc *testExpressionsSuite) TestColumnNameExprRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"abc", "`abc`"},
+		{"`abc`", "`abc`"},
+		{"`ab``c`", "`ab``c`"},
+		{"sabc.tABC", "`sabc`.`tABC`"},
+		{"dabc.sabc.tabc", "`dabc`.`sabc`.`tabc`"},
+		{"dabc.`sabc`.tabc", "`dabc`.`sabc`.`tabc`"},
+		{"`dABC`.`sabc`.tabc", "`dABC`.`sabc`.`tabc`"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr
 	}
+	RunNodeRestoreTest(c, testCases, "select %s", extractNodeFunc)
 }
 
-func (tc *testExpressionsSuite) TestExpresionsRestore(c *C) {
-	parser := parser.New()
-	var testNodes []exprTestCase
-	testNodes = append(testNodes, tc.createTestCase4UnaryOperationExpr()...)
-	testNodes = append(testNodes, tc.createTestCase4ColumnNameExpr()...)
-	testNodes = append(testNodes, tc.createTestCase4IsNullExpr()...)
-
-	for _, node := range testNodes {
-		stmt, err := parser.ParseOneStmt(node.sourceSQL, "", "")
-		comment := Commentf("source %#v", node)
-		c.Assert(err, IsNil, comment)
-		var sb strings.Builder
-		sb.WriteString("SELECT ")
-		err = stmt.(*SelectStmt).Fields.Fields[0].Expr.Restore(&sb)
-		c.Assert(err, IsNil, comment)
-		restoreSql := sb.String()
-		comment = Commentf("source %#v; restore %v", node, restoreSql)
-		c.Assert(restoreSql, Equals, node.expectSQL, comment)
-		stmt2, err := parser.ParseOneStmt(restoreSql, "", "")
-		c.Assert(err, IsNil, comment)
-		CleanNodeText(stmt)
-		CleanNodeText(stmt2)
-		c.Assert(stmt2, DeepEquals, stmt, comment)
+func (tc *testExpressionsSuite) TestIsNullExprRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"a is null", "`a` IS NULL"},
+		{"a is not null", "`a` IS NOT NULL"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr
 	}
+	RunNodeRestoreTest(c, testCases, "select %s", extractNodeFunc)
 }
diff --git a/parser/ast/util_test.go b/parser/ast/util_test.go
@@ -11,12 +11,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package ast
+package ast_test
 
 import (
+	"fmt"
 	"strings"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/parser"
+	. "github.com/pingcap/parser/ast"
 )
 
 var _ = Suite(&testCacheableSuite{})
@@ -82,3 +85,30 @@ func (checker *nodeTextCleaner) Enter(in Node) (out Node, skipChildren bool) {
 func (checker *nodeTextCleaner) Leave(in Node) (out Node, ok bool) {
 	return in, true
 }
+
+type NodeRestoreTestCase struct {
+	sourceSQL string
+	expectSQL string
+}
+
+func RunNodeRestoreTest(c *C, nodeTestCases []NodeRestoreTestCase, template string, extractNodeFunc func(node Node) Node) {
+	parser := parser.New()
+	for _, testCase := range nodeTestCases {
+		sourceSQL := fmt.Sprintf(template, testCase.sourceSQL)
+		expectSQL := fmt.Sprintf(template, testCase.expectSQL)
+		stmt, err := parser.ParseOneStmt(sourceSQL, "", "")
+		comment := Commentf("source %#v", testCase)
+		c.Assert(err, IsNil, comment)
+		var sb strings.Builder
+		err = extractNodeFunc(stmt).Restore(&sb)
+		c.Assert(err, IsNil, comment)
+		restoreSql := fmt.Sprintf(template, sb.String())
+		comment = Commentf("source %#v; restore %v", testCase, restoreSql)
+		c.Assert(restoreSql, Equals, expectSQL, comment)
+		stmt2, err := parser.ParseOneStmt(restoreSql, "", "")
+		c.Assert(err, IsNil, comment)
+		CleanNodeText(stmt)
+		CleanNodeText(stmt2)
+		c.Assert(stmt2, DeepEquals, stmt, comment)
+	}
+}