add Match for mockdb

Author: smallnest

sqlmock/sqlmock.go

@@ -7,6 +7,7 @@ import (
 	"database/sql/driver"
 	"encoding/json"
 	"fmt"
+	"regexp"
 	"sync"
 )
 
@@ -23,6 +24,7 @@ var (
 // ExpectedQuery 表示一个预期的查询
 type ExpectedQuery struct {
 	query   string
+	matcher *regexp.Regexp
 	args    []driver.Value
 	rows    [][]driver.Value
 	columns []string
@@ -55,6 +57,19 @@ func (m *MockDB) ExpectQuery(query string, args ...driver.Value) *ExpectedQuery
 	return eq
 }
 
+// ExpectQuery 期望一个特定的查询
+func (m *MockDB) Macth(matcher string, args ...driver.Value) *ExpectedQuery {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	eq := &ExpectedQuery{
+		args:    args,
+		matcher: regexp.MustCompile(matcher),
+	}
+	m.expected = append(m.expected, eq)
+	return eq
+}
+
 func (eq *ExpectedQuery) WithArgs(args ...driver.Value) *ExpectedQuery {
 	eq.args = args
 	return eq
@@ -141,6 +156,16 @@ func (ms *MockStmt) Exec(args []driver.Value) (driver.Result, error) {
 	defer ms.mockDB.mu.Unlock()
 
 	for i, expected := range ms.mockDB.expected {
+		if expected.matcher != nil && expected.matcher.MatchString(ms.query) && matchArgs(ms.mockDB.expected[i].args, args) {
+			ms.mockDB.expected = append(ms.mockDB.expected[:i], ms.mockDB.expected[i+1:]...)
+
+			if expected.err != nil {
+				return nil, expected.err
+			}
+
+			return &MockResult{}, nil
+		}
+
 		if CompareSQL(expected.query, ms.query) && matchArgs(expected.args, args) {
 			ms.mockDB.expected = append(ms.mockDB.expected[:i], ms.mockDB.expected[i+1:]...)
 
@@ -160,6 +185,16 @@ func (ms *MockStmt) Query(args []driver.Value) (driver.Rows, error) {
 	defer ms.mockDB.mu.Unlock()
 
 	for i, expected := range ms.mockDB.expected {
+		if expected.matcher != nil && expected.matcher.MatchString(ms.query) && matchArgs(ms.mockDB.expected[i].args, args) {
+			ms.mockDB.expected = append(ms.mockDB.expected[:i], ms.mockDB.expected[i+1:]...)
+
+			if expected.err != nil {
+				return nil, expected.err
+			}
+
+			return &MockRows{columns: expected.columns, rows: expected.rows}, nil
+		}
+
 		if CompareSQL(expected.query, ms.query) && matchArgs(expected.args, args) {
 			ms.mockDB.expected = append(ms.mockDB.expected[:i], ms.mockDB.expected[i+1:]...)
 

sqlmock/sqlmock_test.go

@@ -213,3 +213,64 @@ func TestMultipleQueries(t *testing.T) {
 		t.Errorf("unexpected user: %+v", user2)
 	}
 }
+func TestMockDB_Match(t *testing.T) {
+	// 1. 创建 mock 数据库
+	mockDB := sqlmock.NewMock()
+
+	// 2. 期望一个匹配查询并设置返回值
+	pattern := `^SELECT\s+[a-zA-Z0-9_,\s]+FROM\s+[a-zA-Z0-9_]+\s+WHERE\s+[a-zA-Z0-9_]+\s*=\s*\?$`
+
+	mockDB.Macth(pattern).
+		WithArgs("Alice").
+		WillReturnRows([]string{"id", "name", "age"}, [][]driver.Value{
+			{2, "Alice", 25},
+		})
+
+	// 3. 打开数据库连接
+	db, err := mockDB.Open("mock")
+	if err != nil {
+		t.Fatalf("failed to open database: %v", err)
+	}
+
+	// 4. 创建仓库实例
+	repo := &UserRepository{db: db}
+
+	// 5. 执行测试
+	var user User
+	err = repo.db.QueryRow("SELECT id, name, age FROM users WHERE name = ?", "Alice").Scan(&user.ID, &user.Name, &user.Age)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	// 6. 验证结果
+	if user.ID != 2 || user.Name != "Alice" || user.Age != 25 {
+		t.Errorf("unexpected user: %+v", user)
+	}
+}
+
+func TestMockDB_Match_NoMatch(t *testing.T) {
+	// 1. 创建 mock 数据库
+	mockDB := sqlmock.NewMock()
+
+	// 2. 期望一个匹配查询并设置返回值
+	mockDB.Macth("SELECT id, name, age FROM users WHERE name = ?", "Alice").
+		WillReturnRows([]string{"id", "name", "age"}, [][]driver.Value{
+			{2, "Alice", 25},
+		})
+
+	// 3. 打开数据库连接
+	db, err := mockDB.Open("mock")
+	if err != nil {
+		t.Fatalf("failed to open database: %v", err)
+	}
+
+	// 4. 创建仓库实例
+	repo := &UserRepository{db: db}
+
+	// 5. 执行测试
+	var user User
+	err = repo.db.QueryRow("SELECT id, name, age FROM users WHERE name = ?", "Bob").Scan(&user.ID, &user.Name, &user.Age)
+	if err == nil {
+		t.Fatalf("expected error, got nil")
+	}
+}