package db import ( "strings" "testing" ) func TestSplitStatements(t *testing.T) { cases := []struct { in string want []string }{ {"SELECT 1", []string{"SELECT 1"}}, {"SELECT 1; SELECT 2;", []string{"SELECT 1", "SELECT 2"}}, {"SELECT ';'; SELECT 2", []string{"SELECT ';'", "SELECT 2"}}, {"SELECT 1; -- c;mt\nSELECT 2", []string{"SELECT 1", "-- c;mt\nSELECT 2"}}, {"/* a;b */ SELECT 1; SELECT 2", []string{"/* a;b */ SELECT 1", "SELECT 2"}}, {`SELECT "a;b"; SELECT 2`, []string{`SELECT "a;b"`, "SELECT 2"}}, {";;", nil}, {"", nil}, } for _, tc := range cases { got := SplitStatements(tc.in) if len(got) != len(tc.want) { t.Fatalf("SplitStatements(%q) = %#v, want %#v", tc.in, got, tc.want) } for i := range got { if got[i] != tc.want[i] { t.Fatalf("SplitStatements(%q)[%d] = %q, want %q", tc.in, i, got[i], tc.want[i]) } } } } func TestStripComments(t *testing.T) { cases := []struct { in string has string // 结果应包含(大小写敏感原文) missing string // 结果不应包含 }{ {"SELECT 1 -- WHERE x", "SELECT 1", "WHERE x"}, {"UPDATE t SET x=1 # note\nWHERE id=1", "WHERE id=1", "note"}, {"SELECT /* hint */ 1", "SELECT", "hint"}, {"SELECT /*!50000 DROP TABLE t*/", "DROP TABLE t", "/*!"}, {"SELECT 'DROP TABLE x'", "SELECT", "DROP TABLE x"}, {`SELECT "col;drop" FROM t`, "FROM t", "col;drop"}, } for _, tc := range cases { got := StripComments(tc.in) if tc.has != "" && !strings.Contains(got, tc.has) { t.Fatalf("StripComments(%q) = %q, 应包含 %q", tc.in, got, tc.has) } if tc.missing != "" && strings.Contains(got, tc.missing) { t.Fatalf("StripComments(%q) = %q, 不应包含 %q", tc.in, got, tc.missing) } } }