tidb-lite: used for unit testing of Golang database-related code

Use tidb-lite to run TiDB in mocktikv mode in Golang code.

tidb-lite can be used for unit testing of database-related code. If your application uses MySQL/TiDB, there will be a lot of database-related code to be unit tested.

In addition, some applications hope to persist some data locally and store it in a database to facilitate data management. Tidb-lite can also be used in this scenario. Developers use SQL compatible with MySQL protocol to process data locally.

This article mainly introduces how to use tidb-lite for unit testing of database-related codes. So how to write unit tests for database-related code?

The simple way is to start a database instance during the unit test, but this is not very elegant, because the unit test has certain requirements on the environment.

Another common method is to mock SQL services in the test. At present, the more popular solution is to use go-sqlmock.

The problem with go-sqlmock

First, let's take a look at how to use go-sqlmock for unit testing of database-related code.

For example, we have the following codes:

package main

import (
    "database/sql"

    _ "github.com/go-sql-driver/mysql"
)

func recordStats(db *sql.DB, userID, productID int64) ( err error) {
    tx, err := db.Begin()
    if err != nil {
        return
    }

    defer func() {
        switch err {
        case nil:
            err = tx.Commit()
        default:
            tx.Rollback()
        }
    }( )

    if _, err = tx.Exec("UPDATE products SET views = views + 1"); err != nil {
        return
    }
    if _, err = tx.Exec("INSERT INTO product_viewers (user_id, product_id) VALUES (? , ?)", userID, productID); err != nil {
        return
    }
    return
}

func main() {
    // @NOTE: the real connection is not required for tests
    db, err := sql.Open("mysql", "root@/blog")
    if err != nil {
        panic(err)
    }
    defer db.Close()

    if err = recordStats(db, 1 /*some user id*/, 5 /*some product id*/); err != nil {
        panic(err)
    }
}

The recordStats function will increase the number of product views by one and add the user to the list of viewers of the product. The code for unit testing this function using go-sqlmock is as follows:

package main

import (
    "fmt"
    "testing"

    "github.com/DATA-DOG/go-sqlmock"
)

// a successful case
func TestShouldUpdateStats( t *testing.T) {
    db, mock, err := sqlmock.New()
    if err != nil {
        t.Fatalf("an error'%s' was not expected when opening a stub database connection", err)
    }
    defer db.Close()

    mock.ExpectBegin()
    mock.ExpectExec("UPDATE products").WillReturnResult(sqlmock.NewResult(1, 1))
    mock.ExpectExec("INSERT INTO product_viewers").WithArgs(2, 3). WillReturnResult(sqlmock.NewResult(1, 1))
    mock.ExpectCommit()

    // now we execute our method
    if err = recordStats(db, 2, 3); err != nil {
        t.Errorf("error was not expected while updating stats: %s", err)
    }

    // we make sure that all expectations were met
    if err := mock.ExpectationsWereMet(); err != nil {
        t.Errorf("there were unfulfilled expectations: %s", err )
    }
}

go-sqlmock needs to define each step of the database operation and sequence in advance (including the begin and commit of executing the transaction). If the actual operation or steps are inconsistent, an error will be reported. And it is necessary to define the expected return data. If the table is more complicated, it is very troublesome.

In fact, our unit test may only need to pay attention to whether the return result of this function is correct, not the order of execution of each operation of the function.

The advantages of tidb-lite

You can use tidb-lite instead of go-sqlmock.

Simplicity

One of the most important advantages of tidb-lite is simplicity. Run a TiDB directly in the code instead of running a MySQL/TiDB instance before running the unit test. This ensures that the unit test does not depend on the external environment; in addition, we don’t need to write a lot of redundancy and boring test code like go-sqlmock but focus on the correctness of the function.

Compatible with MySQL protocol

TiDB is highly compatible with MySQL protocol. Using tidb-lite can almost completely simulate the MySQL environment.

The usage of tidb-lite

For example, we have the following codes:

package example

import (
    "context"
    "database/sql"
    "fmt"

    "github.com/pingcap/errors"
    "github.com/pingcap/log "
    "go.uber.org/zap"
)

// GetRowCount returns row count of the table.
// if not specify where condition, return total row count of the table.
func GetRowCount(ctx context.Context, db *sql.DB , schemaName string, tableName string, where string) (int64, error) {
    /*
        select count example result:
        mysql> SELECT count(1) cnt from `test`.`itest` where id> 0;
        +----- -+
        | cnt |
        +------+
        | 100 |
        +------+
    */

    query := fmt.Sprintf("SELECT COUNT(1) cnt FROM `%s`.`%s` ", schemaName, tableName)
    if len(where)> 0 {
        query += fmt.Sprintf(" WHERE %s", where)
    }
    log.Debug("get row count", zap.String("sql", query) )

    var cnt sql.NullInt64
    err := db.QueryRowContext(ctx, query).Scan(&cnt)
    if err != nil {
        return 0, errors.Trace(err)
    }
    if !cnt.Valid {
        return 0, errors.NotFoundf ("table `%s`.`%s`", schemaName, tableName)
    }

    return cnt.Int64, nil
}

GetRowCount is used to get the number of eligible rows in the table, the unit test code of this function using tidb-lite is as follows:

package example

import (
    "context"
    "testing"
    "time"

    tidblite "github.com/WangXiangUSTC/tidb-lite"
    . "Github.com/pingcap /check"
)

func TestClient(t *testing.T) {
    TestingT(t)
}

var _ = Suite(&testExampleSuite{})

type testExampleSuite struct{}

func (t *testExampleSuite) TestGetRowCount(c *C) {
    tidbServer, err: = tidblite.NewTiDBServer(tidblite.NewOptions(c.MkDir()))
    c.Assert(err, IsNil)

    dbConn, err := tidbServer.CreateConn()
    c.Assert(err, IsNil)

    ctx, cancel := context.WithTimeout (context.Background(), 10*time.Second)
    defer cancel()

    _, err = dbConn.ExecContext(ctx, "create database example_test")
    c.Assert(err, IsNil)
    _, err = dbConn.ExecContext(ctx , "create table example_test.t(id int primary key, name varchar(24))")
    c.Assert(err, IsNil)
    _, err = dbConn.ExecContext(ctx, "insert into example_test.t values(1, ' a'),(2,'b'),(3,'c')")
    c.Assert(err, IsNil)

    count, err := GetRowCount(ctx, dbConn, "example_test", "t", "id > 2")
    c.Assert(err, IsNil)
    c.Assert(count , Equals, int64(1))

    count, err = GetRowCount(ctx, dbConn, "example_test", "t", "")
    c.Assert(err, IsNil)
    c.Assert(count, Equals, int64(3))
    tidbServer.Close()

    tidbServer2, err := tidblite.NewTiDBServer(tidblite.NewOptions(c.MkDir()))
    c.Assert(err, IsNil)
    defer tidbServer2.Close()

    dbConn2, err := tidbServer2.CreateConn()
    c .Assert(err, IsNil)
    _, err = dbConn2.ExecContext(ctx, "create database example_test")
    c.Assert(err, IsNil)
}

First, we use NewTiDBServer to create a TiDB instance and use CreateConn to get a link to this database. And then you can use this link to access the database, generate test data, and verify the correctness of the function. You can see README for detailed usage.