beego/client/orm/mock/mock_orm.go

// 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, "QueryTable"), []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)
}

// MockRollbackUnlessCommit support RollbackUnlessCommit
func MockRollbackUnlessCommit(err error) *Mock {
	return NewMock(NewSimpleCondition("", "RollbackUnlessCommit"), []interface{}{err}, nil)
}