Support ORM mock

client/orm/mock/mock_orm.go

@@ -0,0 +1,162 @@
+// Copyright 2020 beego
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package mock
+
+import (
+	"database/sql"
+	"os"
+	"path/filepath"
+
+	_ "github.com/mattn/go-sqlite3"
+
+	"github.com/beego/beego/v2/client/orm"
+)
+
+func init() {
+	RegisterMockDB("default")
+}
+
+// RegisterMockDB create an "virtual DB" by using sqllite
+// you should not
+func RegisterMockDB(name string) {
+	source := filepath.Join(os.TempDir(), name+".db")
+	_ = orm.RegisterDataBase(name, "sqlite3", source)
+}
+
+// MockTable only check table name
+func MockTable(tableName string, resp ...interface{}) *Mock {
+	return NewMock(NewSimpleCondition(tableName, ""), resp, nil)
+}
+
+// MockMethod only check method name
+func MockMethod(method string, resp ...interface{}) *Mock {
+	return NewMock(NewSimpleCondition("", method), resp, nil)
+}
+
+// MockOrmRead support orm.Read and orm.ReadWithCtx
+// cb is used to mock read data from DB
+func MockRead(tableName string, cb func(data interface{}), err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "ReadWithCtx"), []interface{}{err}, func(inv *orm.Invocation) {
+		if cb != nil {
+			cb(inv.Args[0])
+		}
+	})
+}
+
+// MockReadForUpdateWithCtx support ReadForUpdate and ReadForUpdateWithCtx
+// cb is used to mock read data from DB
+func MockReadForUpdateWithCtx(tableName string, cb func(data interface{}), err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "ReadForUpdateWithCtx"),
+		[]interface{}{err},
+		func(inv *orm.Invocation) {
+			cb(inv.Args[0])
+		})
+}
+
+// MockReadOrCreateWithCtx support ReadOrCreate and ReadOrCreateWithCtx
+// cb is used to mock read data from DB
+func MockReadOrCreateWithCtx(tableName string,
+	cb func(data interface{}),
+	insert bool, id int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "ReadOrCreateWithCtx"),
+		[]interface{}{insert, id, err},
+		func(inv *orm.Invocation) {
+			cb(inv.Args[0])
+		})
+}
+
+// MockInsertWithCtx support Insert and InsertWithCtx
+func MockInsertWithCtx(tableName string, id int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "InsertWithCtx"), []interface{}{id, err}, nil)
+}
+
+// MockInsertMultiWithCtx support InsertMulti and InsertMultiWithCtx
+func MockInsertMultiWithCtx(tableName string, cnt int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "InsertMultiWithCtx"), []interface{}{cnt, err}, nil)
+}
+
+// MockInsertOrUpdateWithCtx support InsertOrUpdate and InsertOrUpdateWithCtx
+func MockInsertOrUpdateWithCtx(tableName string, id int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "InsertOrUpdateWithCtx"), []interface{}{id, err}, nil)
+}
+
+// MockUpdateWithCtx support UpdateWithCtx and Update
+func MockUpdateWithCtx(tableName string, affectedRow int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "UpdateWithCtx"), []interface{}{affectedRow, err}, nil)
+}
+
+// MockDeleteWithCtx support Delete and DeleteWithCtx
+func MockDeleteWithCtx(tableName string, affectedRow int64, err error) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "DeleteWithCtx"), []interface{}{affectedRow, err}, nil)
+}
+
+// MockQueryM2MWithCtx support QueryM2MWithCtx and QueryM2M
+// Now you may be need to use golang/mock to generate QueryM2M mock instance
+// Or use DoNothingQueryM2Mer
+// for example:
+// 	post := Post{Id: 4}
+// 	m2m := Ormer.QueryM2M(&post, "Tags")
+// when you write test code:
+// MockQueryM2MWithCtx("post", "Tags", mockM2Mer)
+// "post" is the table name of model Post structure
+// TODO provide orm.QueryM2Mer
+func MockQueryM2MWithCtx(tableName string, name string, res orm.QueryM2Mer) *Mock {
+	return NewMock(NewQueryM2MerCondition(tableName, name), []interface{}{res}, nil)
+}
+
+// MockLoadRelatedWithCtx support LoadRelatedWithCtx and LoadRelated
+func MockLoadRelatedWithCtx(tableName string, name string, rows int64, err error) *Mock {
+	return NewMock(NewQueryM2MerCondition(tableName, name), []interface{}{rows, err}, nil)
+}
+
+// MockQueryTableWithCtx support QueryTableWithCtx and QueryTable
+func MockQueryTableWithCtx(tableName string, qs orm.QuerySeter) *Mock {
+	return NewMock(NewSimpleCondition(tableName, "QueryTableWithCtx"), []interface{}{qs}, nil)
+}
+
+// MockRawWithCtx support RawWithCtx and Raw
+func MockRawWithCtx(rs orm.RawSeter) *Mock {
+	return NewMock(NewSimpleCondition("", "RawWithCtx"), []interface{}{rs}, nil)
+}
+
+// MockDriver support Driver
+// func MockDriver(driver orm.Driver) *Mock {
+// 	return NewMock(NewSimpleCondition("", "Driver"), []interface{}{driver})
+// }
+
+// MockDBStats support DBStats
+func MockDBStats(stats *sql.DBStats) *Mock {
+	return NewMock(NewSimpleCondition("", "DBStats"), []interface{}{stats}, nil)
+}
+
+// MockBeginWithCtxAndOpts support Begin, BeginWithCtx, BeginWithOpts, BeginWithCtxAndOpts
+// func MockBeginWithCtxAndOpts(txOrm *orm.TxOrmer, err error) *Mock {
+// 	return NewMock(NewSimpleCondition("", "BeginWithCtxAndOpts"), []interface{}{txOrm, err})
+// }
+
+// MockDoTxWithCtxAndOpts support DoTx, DoTxWithCtx, DoTxWithOpts, DoTxWithCtxAndOpts
+// func MockDoTxWithCtxAndOpts(txOrm *orm.TxOrmer, err error) *Mock {
+// 	return MockBeginWithCtxAndOpts(txOrm, err)
+// }
+
+// MockCommit support Commit
+func MockCommit(err error) *Mock {
+	return NewMock(NewSimpleCondition("", "Commit"), []interface{}{err}, nil)
+}
+
+// MockRollback support Rollback
+func MockRollback(err error) *Mock {
+	return NewMock(NewSimpleCondition("", "Rollback"), []interface{}{err}, nil)
+}