如何为使用gocql的Go函数编写无真实数据库依赖的单元测试

要在不连接真实Cassandra的情况下测试GetSeglistFromCassandra，核心是通过接口抽象+Mock实现替换gocql的真实依赖。以下是具体步骤和代码示例：

1. 抽象gocql核心依赖接口 #

原代码直接依赖*gocql.Session、*gocql.Query等具体类型，无法直接Mock。我们先定义对应接口，只保留业务代码用到的方法：

// 定义抽象接口，仅包含业务代码用到的方法
type CassandraSession interface {
    Query(stmt string, values ...interface{}) CassandraQuery
}

type CassandraQuery interface {
    Consistency(c gocql.Consistency) CassandraQuery
    Exec() error
    Iter() CassandraIter
}

type CassandraIter interface {
    Scan(dest ...interface{}) bool
    Close() error
}

2. 改造原业务代码依赖抽象接口 #

修改原代码，让getSelectQuery和GetSeglistFromCassandra依赖上面定义的接口，而不是gocql的具体类型：

// 全局变量改为抽象类型
var CassandraSession CassandraSession

func getSelectQuery(ClientId, UID, UUID, tableName, selectColumns string, sess CassandraSession) (CassandraQuery, error) {
    selectQuery := fmt.Sprintf(qTamplate, selectColumns, tableName)
    query := sess.Query(selectQuery, ClientId, UID).Consistency(gocql.One)
    err := query.Exec()
    if err != nil {
        return query, err
    }
    return query, nil
}

func GetSeglistFromCassandra(reqId string, dataMap MapUserRequest) (models.IsInSegmentResponse, error) {
    resutMap := models.IsInSegmentResponse{}
    // 此处假设cid为dataMap.ClientId的处理逻辑，原代码未定义，测试时可忽略具体实现
    cid := dataMap.ClientId
    tableName := fmt.Sprintf("getter_%v", cid[len(cid)-1])
    cqlsh := CassandraSession
    selectColums := "cohortid, is_deleted, cohortype"
    selectQueryObj, err := getSelectQuery(dataMap.ClientId, dataMap.UID, dataMap.UUID, tableName, selectColums, cqlsh)
    if err != nil {
        return resutMap, err
    }
    var id, deleted, cohortType int
    iter := selectQueryObj.Iter()
    for iter.Scan(&id, &deleted, &cohortType) {
        if deleted == 0 {
            if cohortType == 3 || cohortType == 2 {
                resutMap.ListId = append(resutMap.ListId, id)
            } else {
                resutMap.SegIds = append(resutMap.SegIds, id)
            }
        }
    }
    if err := iter.Close(); err != nil {
        return resutMap, err
    }
    return resutMap, nil
}

3. 实现Mock结构体 #

编写Mock实现，控制每个方法的返回值和行为，满足不同测试场景：

// MockSession 实现CassandraSession接口
type MockSession struct {
    QueryFunc func(stmt string, values ...interface{}) CassandraQuery
}

func (m *MockSession) Query(stmt string, values ...interface{}) CassandraQuery {
    return m.QueryFunc(stmt, values...)
}

// MockQuery 实现CassandraQuery接口
type MockQuery struct {
    ConsistencyFunc func(c gocql.Consistency) CassandraQuery
    ExecFunc        func() error
    IterFunc        func() CassandraIter
}

func (m *MockQuery) Consistency(c gocql.Consistency) CassandraQuery {
    return m.ConsistencyFunc(c)
}

func (m *MockQuery) Exec() error {
    return m.ExecFunc()
}

func (m *MockQuery) Iter() CassandraIter {
    return m.IterFunc()
}

// MockIter 实现CassandraIter接口
type MockIter struct {
    ScanFunc  func(dest ...interface{}) bool
    CloseFunc func() error
}

func (m *MockIter) Scan(dest ...interface{}) bool {
    return m.ScanFunc(dest...)
}

func (m *MockIter) Close() error {
    return m.CloseFunc()
}

4. 编写单元测试用例 #

针对不同场景编写测试，比如正常返回数据、查询失败、迭代关闭失败等：

package main

import (
    "testing"
    "github.com/gocql/gocql"
    "github.com/stretchr/testify/assert"
)

// 测试正常返回数据的场景
func TestGetSeglistFromCassandra_Success(t *testing.T) {
    // 构造Mock迭代器，模拟返回两条数据
    scanCount := 0
    mockIter := &MockIter{
        ScanFunc: func(dest ...interface{}) bool {
            scanCount++
            switch scanCount {
            case 1:
                // 第一条数据：id=1, deleted=0, cohortType=1 → 加入SegIds
                *dest[0].(*int) = 1
                *dest[1].(*int) = 0
                *dest[2].(*int) = 1
                return true
            case 2:
                // 第二条数据：id=2, deleted=0, cohortType=2 → 加入ListId
                *dest[0].(*int) = 2
                *dest[1].(*int) = 0
                *dest[2].(*int) = 2
                return true
            default:
                return false
            }
        },
        CloseFunc: func() error {
            return nil
        },
    }

    mockQuery := &MockQuery{
        ConsistencyFunc: func(c gocql.Consistency) CassandraQuery {
            return mockQuery
        },
        ExecFunc: func() error {
            return nil
        },
        IterFunc: func() CassandraIter {
            return mockIter
        },
    }

    mockSession := &MockSession{
        QueryFunc: func(stmt string, values ...interface{}) CassandraQuery {
            return mockQuery
        },
    }

    // 替换全局Session为Mock
    CassandraSession = mockSession

    // 构造测试请求
    dataMap := MapUserRequest{
        ClientId: "123",
        UID:      "uid1",
        UUID:     "uuid1",
    }

    resp, err := GetSeglistFromCassandra("req1", dataMap)
    assert.NoError(t, err)
    assert.Equal(t, []int{1}, resp.SegIds)
    assert.Equal(t, []int{2}, resp.ListId)
}

// 测试查询执行失败的场景
func TestGetSeglistFromCassandra_QueryExecError(t *testing.T) {
    mockQuery := &MockQuery{
        ConsistencyFunc: func(c gocql.Consistency) CassandraQuery {
            return mockQuery
        },
        ExecFunc: func() error {
            return assert.AnError
        },
    }

    mockSession := &MockSession{
        QueryFunc: func(stmt string, values ...interface{}) CassandraQuery {
            return mockQuery
        },
    }

    CassandraSession = mockSession

    dataMap := MapUserRequest{
        ClientId: "123",
        UID:      "uid1",
        UUID:     "uuid1",
    }

    resp, err := GetSeglistFromCassandra("req1", dataMap)
    assert.Error(t, err)
    assert.Equal(t, models.IsInSegmentResponse{}, resp)
}

// 测试迭代关闭失败的场景
func TestGetSeglistFromCassandra_IterCloseError(t *testing.T) {
    mockIter := &MockIter{
        ScanFunc: func(dest ...interface{}) bool {
            return false // 无数据返回
        },
        CloseFunc: func() error {
            return assert.AnError
        },
    }

    mockQuery := &MockQuery{
        ConsistencyFunc: func(c gocql.Consistency) CassandraQuery {
            return mockQuery
        },
        ExecFunc: func() error {
            return nil
        },
        IterFunc: func() CassandraIter {
            return mockIter
        },
    }

    mockSession := &MockSession{
        QueryFunc: func(stmt string, values ...interface{}) CassandraQuery {
            return mockQuery
        },
    }

    CassandraSession = mockSession

    dataMap := MapUserRequest{
        ClientId: "123",
        UID:      "uid1",
        UUID:     "uuid1",
    }

    resp, err := GetSeglistFromCassandra("req1", dataMap)
    assert.Error(t, err)
}

关键说明 #

• 接口抽象只保留业务用到的方法，避免过度抽象；
• Mock结构体通过函数字段灵活控制每个方法的行为，适配不同测试场景；
• 测试前替换全局的CassandraSession为Mock实现，彻底隔离真实依赖；
• 借助testify/assert库简化断言逻辑，提升测试可读性。

内容的提问来源于stack exchange，提问作者Mesc