單元測試中如何解決 MySQL 存儲依賴問題

在編寫單元測試的過程中，如果被測試代碼有外部依賴，爲了便於測試，我們就要想辦法來解決這些外部依賴問題。在做 Web 開發時，MySQL 存儲就是一個非常常見的外部依賴，本文就來探討在 Go 語言中編寫單元測試時，如何解決 MySQL 存儲依賴。

HTTP 服務程序示例

假設我們有一個 HTTP 服務程序對外提供服務，代碼如下：

main.go

package main

import (
 "encoding/json"
 "fmt"
 "gorm.io/gorm"
 "io"
 "net/http"
 "strconv"

 "github.com/julienschmidt/httprouter"

 "github.com/jianghushinian/blog-go-example/test/mysql/store"
)

func NewUserHandler(db *gorm.DB) *UserHandler {
 return &UserHandler{
  store: store.NewUserStore(db),
 }
}

type UserHandler struct {
 store store.UserStore
}

func (h *UserHandler) CreateUser(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
 ...
}

func (h *UserHandler) GetUser(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 ...
}

func setupRouter(handler *UserHandler) *httprouter.Router {
 router := httprouter.New()
 router.POST("/users", handler.CreateUser)
 router.GET("/users/:id", handler.GetUser)
 return router
}

func main() {
 mysqlDB, _ := store.NewMySQLDB("localhost", "3306", "user", "password", "test")
 handler := NewUserHandler(mysqlDB)
 router := setupRouter(handler)

 _ = http.ListenAndServe(":8000", router)
}

這個服務監聽 8000 端口，分別提供了兩個 HTTP 接口：

POST /users 用來創建用戶。

GET /users/:id 用來獲取指定 ID 對應的用戶信息。

UserHandler 是一個結構體，它依賴外部存儲接口 store.UserStore，這個接口定義如下：

store/store.go

package store

import "gorm.io/gorm"

type UserStore interface {
 Create(user *User) error
 Get(id int) (*User, error)
}

func NewUserStore(db *gorm.DB) UserStore {
 return &userStore{db}
}

type userStore struct {
 db *gorm.DB
}

func (s *userStore) Create(user *User) error {
 return s.db.Create(user).Error
}

func (s *userStore) Get(id int) (*User, error) {
 var user User
 err := s.db.First(&user, id).Error
 return &user, err
}

store.UserStore 定義了兩個方法，分別用來創建、獲取用戶信息。

User 模型定義如下：

store/model.go

type User struct {
 ID   int    `gorm:"id"`
 Name string `gorm:"name"`
}

store.userStore 結構體則實現了 store.UserStore 接口。

store.userStore 結構體又依賴了 GORM 庫的 *gorm.DB 類型，表示一個數據庫連接對象。

我們可以使用 NewMySQLDB 建立數據庫連接得到 *gorm.DB 對象：

store/mysql.go

func NewMySQLDB(host, port, user, pass, dbname string) (*gorm.DB, error) {
 dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
  user, pass, host, port, dbname)
 return gorm.Open(mysql.Open(dsn), &gorm.Config{})
}

至此，這個 HTTP 服務程序整體邏輯就基本介紹完了。

其目錄結構如下：

$ tree
.
├── go.mod
├── go.sum
├── main.go
└── store
    ├── model.go
    ├── mysql.go
    └── store.go

爲了保證業務的正確性，我們應該對 (*UserHandler).CreateUser 和 (*UserHandler).GetUser 這兩個 Handler 進行單元測試。

這兩個 Handler 定義如下：

func (h *UserHandler) CreateUser(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
 w.Header().Set("Content-Type", "application/json")

 body, err := io.ReadAll(r.Body)
 if err != nil {
  w.WriteHeader(http.StatusBadRequest)
  _, _ = fmt.Fprintf(w, `{"msg":"%s"}`, err.Error())
  return
 }
 defer func() { _ = r.Body.Close() }()

 u := store.User{}
 if err := json.Unmarshal(body, &u); err != nil {
  w.WriteHeader(http.StatusInternalServerError)
  _, _ = fmt.Fprintf(w, `{"msg":"%s"}`, err.Error())
  return
 }

 if err := h.store.Create(&u); err != nil {
  w.WriteHeader(http.StatusInternalServerError)
  _, _ = fmt.Fprintf(w, `{"msg":"%s"}`, err.Error())
  return
 }
 w.WriteHeader(http.StatusCreated)
}

func (h *UserHandler) GetUser(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 id := ps[0].Value
 uid, _ := strconv.Atoi(id)

 w.Header().Set("Content-Type", "application/json")
 u, err := h.store.Get(uid)
 if err != nil {
  w.WriteHeader(http.StatusInternalServerError)
  _, _ = fmt.Fprintf(w, `{"msg":"%s"}`, err.Error())
  return
 }
 _, _ = fmt.Fprintf(w, `{"id":%d,"name":"%s"}`, u.ID, u.Name)
}

不過，由於文章篇幅所限，我這裏僅以測試 (*UserHandler).GetUser 方法爲例，演示如何在測試過程中解決 MySQL 依賴問題，對 (*UserHandler).CreateUser 方法的測試就當做作業留給你自己來完成了（當然，你也可以到我的 GitHub 上查看我的實現）。

Fake 測試

我們要爲 (*UserHandler).GetUser 方法編寫單元測試，首先就要分析下這個方法的外部依賴。

func (h *UserHandler) GetUser(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
 id := ps[0].Value
 uid, _ := strconv.Atoi(id)

 w.Header().Set("Content-Type", "application/json")
 u, err := h.store.Get(uid)
 if err != nil {
  w.WriteHeader(http.StatusInternalServerError)
  _, _ = fmt.Fprintf(w, `{"msg":"%s"}`, err.Error())
  return
 }
 _, _ = fmt.Fprintf(w, `{"id":%d,"name":"%s"}`, u.ID, u.Name)
}

UserHandler 結構本身依賴了 store.UserStore，這是一個接口，定義了創建和獲取用戶信息的兩個方法。

我們使用實現了 store.UserStore 接口的 store.userStore 結構體來初始化 UserHandler：

func NewUserHandler(db *gorm.DB) *UserHandler {
 return &UserHandler{
  store: store.NewUserStore(db),
 }
}

func NewUserStore(db *gorm.DB) UserStore {
 return &userStore{db}
}

store.userStore 結構體會使用 GORM 來完成對 MySQL 數據庫的操作。所以，我們分析出 GetUser 方法的第一個外部依賴實際上就是 MySQL 存儲。

GetUser 方法還接收三個參數，它們都屬於 HTTP 網絡相關的外部依賴，你可以在我的另一篇文章《在 Go 語言單元測試中如何解決 HTTP 網絡依賴問題》中找到解決方案，就不在本文中進行講解了。

所以，我們現在重點要關注的就只有一個問題，如何解決 MySQL 存儲依賴。

我們來整理下 MySQL 外部依賴的程序調用鏈：

MySQL 外部依賴

可以發現，store.UserStore 接口是 UserHandler 和 store.userStore 結構體建立連接的橋樑，我們可以將它作爲突破口，實現一個 Fake object，來替換 store.userStore 結構體。

Fake object

所謂 Fake object，其實就是我們同樣要定義一個結構體，並實現 Create 和 Get 兩個方法，以此來實現 store.UserStore 接口。

type fakeUserStore struct{}

func (f *fakeUserStore) Create(user *store.User) error {
 return nil
}

func (f *fakeUserStore) Get(id int) (*store.User, error) {
 return &store.User{ID: id, Name: "test"}, nil
}

與 store.userStore 結構體不同，fakeUserStore 並不依賴 *gorm.DB，也就不涉及 MySQL 數據庫操作了，這樣就解決了 MySQL 外部存儲依賴。

(*fakeUserStore).Create 方法沒做任何操作，直接返回 nil，(*fakeUserStore).Get 方法則根據傳進來的 id 返回固定的 User 信息。這也是 Fake object 的特點，爲真實對象實現一個簡化版本。

這樣，我們在編寫測試代碼時，只需要取代 store.userStore 結構體，使用 fakeUserStore 來實例化 UserHandler，就可以避免與 MySQL 數據庫打交道了。

handler := &UserHandler{store: &fakeUserStore{}}

爲 (*UserHandler).GetUser 方法編寫的單元測試完整代碼如下：

func TestUserHandler_GetUser_by_fake(t *testing.T) {
 handler := &UserHandler{store: &fakeUserStore{}}
 router := setupRouter(handler)

 w := httptest.NewRecorder()
 req := httptest.NewRequest("GET", "/users/1", nil)
 router.ServeHTTP(w, req)

 assert.Equal(t, 200, w.Code)
 assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
 assert.Equal(t, `{"id":1,"name":"test"}`, w.Body.String())
}

現在被測試的 (*UserHandler).GetUser 方法中通過 h.store.Get(uid) 從數據庫中獲取用戶信息時，就不用再去查詢 MySQL 了，而是由 (*fakeUserStore).Get 方法直接返回 Fake 數據。

使用 go test 來執行測試函數：

$ go test -v -run="TestUserHandler_GetUser_by_fake"
=== RUN   TestUserHandler_GetUser_by_fake
--- PASS: TestUserHandler_GetUser_by_fake (0.00s)
PASS
ok      github.com/jianghushinian/blog-go-example/test/mysql    0.465s

測試通過。

可以發現，使用 Fake 測試來解決 MySQL 外部依賴還是比較簡單的，我們僅需要參考 store.userStore 實現一個簡化版本的 fakeUserStore，然後在測試過程中，使用簡化版本的 fakeUserStore 對象替換掉 store.userStore 即可。

Mock 測試

前文中，我們使用 fakeUserStore 來替換 store.userStore，以此來接口 MySQL 依賴問題。

不過，這種使用 Fake object 來解決外部依賴的方式存在兩個較爲常見的弊端：

一個是使用 Fake object 需要手動編寫大量代碼，這裏的 store.UserStore 接口僅定義了兩個方法還好，但一個線上的複雜業務，可能有幾十個接口，每個接口又有幾十個方法，此時如果還是手動來編寫這些代碼，需要消耗大量時間。

另一個是 Fake object 返回結果比較固定，如果想測試其他情況，比如查詢的 User 不存在，需要報錯的情況，就得在 (*fakeUserStore).Get 方法中編寫更多的邏輯，這增加了實現 Fake object 的複雜度。

那麼有沒有一種替代方案，來彌補 Fake object 的這兩個弊端呢？

答案是使用 Mock 測試。

Mock 和 Fake 類似，本質上都是使用一個對象，去替代另一個對象。Fake 測試是實現了一個真實對象（store.userStore）的簡化版本（fakeUserStore），Mock 測試則是使用模擬對象來斷言真實對象被調用時的輸入符合預期，然後通過模擬對象返回指定輸出。

在 Go 中，我們可以使用 gomock 來實現 Mock 測試。

gomock 項目起源於 Google 的 golang/mock 倉庫。不幸的是，谷歌不再維護這個項目了。幸運的是，這個項目由 Uber fork 了一份，並繼續維護。

gomock 包含兩個部分：gomock 包和 mockgen 命令行工具。gomock 包用來完成對被 Mock 對象的生命週期管理，mockgen 工具則用來自動生成 Mock 代碼。

可以通過如下方式來安裝 gomock 包和 mockgen 工具：

$ go get go.uber.org/mock/gomock@latest
$ go install go.uber.org/mock/mockgen@latest

注意：在項目根目錄下通過 go get 命令獲取 gomock 包後，不要急着執行 go mod tidy，因爲現在 gomock 包屬於 indirect 依賴，還沒有被使用。當通過 mockgen 工具生成了 Mock 代碼以後，再來執行 go mod tidy，go.mod 文件中才不會丟失 gomock 依賴。

要想使用 gomock 來模擬 store.UserStore 接口的實現，我們先要使用 mockgen 工具來生成 Mock 代碼：

 $ mockgen -source store/store.go -destination store/mocks/gomock.go -package mocks

-source 參數指明需要 Mock 的接口文件路徑，即 store.UserStore 接口所在文件。

-destination 參數指明生成的 Mock 文件路徑。

-package 參數指明生成的 Mock 文件包名。

在項目根目錄下執行 mockgen 命令，即可生成 Mock 文件：

// Code generated by MockGen. DO NOT EDIT.
// Source: store/store.go

// Package mocks is a generated GoMock package.
package mocks

import (
 reflect "reflect"

 store "github.com/jianghushinian/blog-go-example/test/mysql/store"
 gomock "go.uber.org/mock/gomock"
)

// MockUserStore is a mock of UserStore interface.
type MockUserStore struct {
 ctrl     *gomock.Controller
 recorder *MockUserStoreMockRecorder
}

// MockUserStoreMockRecorder is the mock recorder for MockUserStore.
type MockUserStoreMockRecorder struct {
 mock *MockUserStore
}

// NewMockUserStore creates a new mock instance.
func NewMockUserStore(ctrl *gomock.Controller) *MockUserStore {
 mock := &MockUserStore{ctrl: ctrl}
 mock.recorder = &MockUserStoreMockRecorder{mock}
 return mock
}

// EXPECT returns an object that allows the caller to indicate expected use.
func (m *MockUserStore) EXPECT() *MockUserStoreMockRecorder {
 return m.recorder
}

// Create mocks base method.
func (m *MockUserStore) Create(user *store.User) error {
 m.ctrl.T.Helper()
 ret := m.ctrl.Call(m, "Create", user)
 ret0, _ := ret[0].(error)
 return ret0
}

// Create indicates an expected call of Create.
func (mr *MockUserStoreMockRecorder) Create(user interface{}) *gomock.Call {
 mr.mock.ctrl.T.Helper()
 return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockUserStore)(nil).Create), user)
}

// Get mocks base method.
func (m *MockUserStore) Get(id int) (*store.User, error) {
 m.ctrl.T.Helper()
 ret := m.ctrl.Call(m, "Get", id)
 ret0, _ := ret[0].(*store.User)
 ret1, _ := ret[1].(error)
 return ret0, ret1
}

// Get indicates an expected call of Get.
func (mr *MockUserStoreMockRecorder) Get(id interface{}) *gomock.Call {
 mr.mock.ctrl.T.Helper()
 return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockUserStore)(nil).Get), id)
}

提示：生成的 mocks 包代碼你無需全部看懂，僅知道它大概生成了什麼內容，如何使用即可。

可以發現，mockgen 爲我們生成了 mocks.MockUserStore 結構體，並且實現了 Create、Get 兩個方法，即實現了 store.UserStore 接口。

現在，我們就可以使用生成的 Mock 對象來編寫單元測試代碼了：

func TestUserHandler_GetUser_by_mock(t *testing.T) {
 ctrl := gomock.NewController(t)
 // 斷言 mockUserStore.Get 方法會被調用
 defer ctrl.Finish()

 mockUserStore := mocks.NewMockUserStore(ctrl)
 mockUserStore.EXPECT().Get(2).Return(&store.User{
  ID:   2,
  Name: "user2",
 }, nil)

 handler := &UserHandler{store: mockUserStore}
 router := setupRouter(handler)

 w := httptest.NewRecorder()
 req := httptest.NewRequest("GET", "/users/2", nil)
 router.ServeHTTP(w, req)

 assert.Equal(t, 200, w.Code)
 assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
 assert.Equal(t, `{"id":2,"name":"user2"}`, w.Body.String())
}

gomock.NewController(t) 用來創建一個 Mock 控制器，該對象可以控制整個 Mock 生命週期。

ctrl.Finish() 用來斷言 Mock 對象使用 EXPECT() 方法設置的期待執行方法會被調用，一般使用 defer 語句來調用，防止最後忘記。不過，如果你使用的 Go 版本大於 1.14，則可以不必顯式調用 ctrl.Finish()。

mocks.NewMockUserStore(ctrl) 使用 Mock 控制器創建了 *mocks.MockUserStore 對象，有了它，我們就可以模擬調用 store.UserStore 接口對應方法的邏輯了：

mockUserStore.EXPECT().Get(2).Return(&store.User{
    ID:   2,
    Name: "user2",
}, nil)

mockUserStore 對象就相當於我們前文中實現的 fakeUserStore。

Mock 對象的 EXPECT() 方法用來設置預期被調用的方法，以及被調用方法所期望的輸入，它支持鏈式調用，.Get(2) 表示期望在測試中調用 Mock 對象 mockUserStore 的 Get 方法時，輸入參數是 2，Return 方法用來設置輸出，即返回值內容。

這就相當於，我們實現了 fakeUserStore 的 Get 方法。

我們可以使用 mockUserStore 來實例化 UserHandler 對象。

在 req 請求中，我們設置請求的用戶 ID 值爲 2，即 mockUserStore 對象斷言中的參數，二者參數匹配，Mock 對象才能生效。

單元測試最後，斷言了返回結果爲 {"id":2,"name":"user2"}，即 mockUserStore 對象期望的返回結果。

現在我們就可以測試 (*UserHandler).GetUser 方法了。

使用 go test 來執行測試函數：

$ go test -v -run="TestUserHandler_GetUser_by_mock"
=== RUN   TestUserHandler_GetUser_by_mock
--- PASS: TestUserHandler_GetUser_by_mock (0.00s)
PASS
ok      github.com/jianghushinian/blog-go-example/test/mysql    0.220s

測試通過。

使用 Mock 測試來解決 MySQL 外部依賴問題，我們無需手動編寫 Mock 對象的代碼，可以使用 mockgen 工具爲我們自動生成，簡化了 Fake 測試中編寫 fakeUserStore 的過程。

並且，如果想要測試其他情況，僅需要再次使用 Mock 對象的 EXPECT() 方法來設置 Get 方法的期望輸入和輸出即可。

比如設置預期查詢 ID 爲 3 的用戶信息時，返回 user not found 錯誤：

mockUserStore.EXPECT().Get(3).Return(nil, errors.New("user not found"))

Mock 測試更方便我們測試不同業務場景。

gomock 更多用法

gomock 還有一些使用技巧值得分享。

mockgen

前文中，我們使用 mockgen 通過指定源碼文件形式生成了 Mock 代碼：

$ mockgen -source store/store.go -destination store/mocks/gomock.go -package mocks

mockgen 工具還支持通過反射模式來生成 Mock 代碼：

$ mockgen -package mocks -destination store/mocks/gomock.go github.com/jianghushinian/blog-go-example/test/mysql/store UserStore

命令最後的兩個參數分別代表需要生成 Mock 代碼的包的導入路徑和逗號分隔的接口列表。

執行以上命令同樣能夠成功生成 Mock 代碼。

此外，我們還可以將 mockgen 命令寫到 Go 文件中，然後使用 Go generate 工具來生成 Mock 代碼：

store/generate.go

package store

//go:generate mockgen -package mocks -destination ./mocks/gomock.go . UserStore

這次我們的 mockgen 命令又有所不同，包的導入路徑僅爲一個 .，表示當前目錄，這也是被支持的。

這時候，我們只需要在項目根目錄下執行 go generate ./... 命令即可生成 Mock 代碼。./... 表示查找項目下全部文件，go generate 會自動找到帶有 //go:generate 註釋的命令並執行。

如果我們有多個源碼文件要生成 Mock 代碼，go generate 方式就非常合適，僅需要在 Go 文件中分多行依次寫出 mockgen 命令即可使用一條命令一次全部生成。

gomock

前文中，我們使用了 Mock 對象 mockUserStore 的 EXPECT() 方法來設置 Get 方法所期待的輸入和輸出。

mockUserStore.EXPECT().Get(2).Return(&store.User{
    ID:   2,
    Name: "user2",
}, nil)

有時候，EXPECT() 所作用的方法可能存在多個參數，且有些參數不容易模擬，比如最常見的 context.Context 參數，針對這些情況，gomock 提供了更多的參數匹配方法：

gomock.Any() 表示匹配任意參數，適合參數模擬困難的情況。

gomock.Eq(x) 表示匹配與 x 相等的參數。

gomock.Not(x) 表示匹配與 x 不想等的參數。

gomock.Nil() 表示匹配 nil 參數。

gomock.Len(i) 表示匹配長度爲 i 的參數。

gomock.All(ms) 表示傳入的所有參數都想等才能匹配。

以上這些參數匹配方法都可以像如下這樣使用：

mockUserStore.EXPECT().Get(gomock.Eq(2)).Return(&store.User{
    ID:   2,
    Name: "user2",
}, nil)

此外，我們可以約束 EXPECT() 所作用方法的執行次數：

.Return(xxx).Times(2) // 預期方法會被調用 2 次
.Return(xxx).MaxTimes(2) // 預期方法最多執行 2 次
.Return(xxx).MinTimes(2) // 預期方法至少執行 2 次
.Return(xxx).AnyTimes() // 預期方法執行任意次都能匹配

還可以約束 EXPECT() 所作用方法的執行順序：

.Return(xxx).After(preReq) // 當前預期方法在 preReq 預期方法執行完成之後執行

以上便是我認爲 gomock 中比較常用的功能講解，更多功能可參考官方文檔。

總結

本文向大家介紹了在 Go 中編寫單元測試時，如何解決 MySQL 外部依賴的問題。

我們分別使用了 Fake object 和 Mock 兩種方式，來替換原有的外部依賴。

Web 服務的代碼不是隨意設計的，有意將 UserHandler 依賴的類型設爲 store.UserStore 接口，而不是 store.userStore 結構體，是爲了解耦。通過使用接口，解決了 UserHandler 與 store.userStore 結構體強綁定的問題，這就給我們使用 fakeUserStore 或 mockUserStore 來替代 store.userStore 創造了機會。

可以發現，本文介紹的兩種方法其實不僅能夠用於解決 MySQL 外部依賴問題。任何使用接口編寫的代碼，在測試時都可以使用這兩種方式來替換依賴。這就是 Go 面向接口編程的好處。

本文完整代碼示例我放在了 GitHub 上，歡迎點擊查看。

希望此文能對你有所幫助。

參考

• gomock 源碼：https://github.com/golang/mock/

• Uber gomock 源碼：https://github.com/uber/mock

• Uber gomock 文檔：https://pkg.go.dev/go.uber.org/mock/gomock

• 本文 GitHub 源碼：https://github.com/jianghushinian/blog-go-example/tree/main/test/mysql

聯繫我

微信：jianghushinian

郵箱：jianghushinian007@outlook.com

博客地址：https://jianghushinian.cn

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/Bn1g_p6vv4icK38VcxJbcg