github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/orm/model/logic_epoch_test.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package model
    15  
    16  import (
    17  	"context"
    18  	"regexp"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/DATA-DOG/go-sqlmock"
    23  	gsql "github.com/go-sql-driver/mysql"
    24  	"github.com/pingcap/failpoint"
    25  	"github.com/pingcap/tiflow/pkg/errors"
    26  	"github.com/stretchr/testify/require"
    27  	"gorm.io/driver/mysql"
    28  	"gorm.io/gorm"
    29  )
    30  
    31  func mockGetDBConn(t *testing.T) (*gorm.DB, sqlmock.Sqlmock) {
    32  	db, mock, err := sqlmock.New()
    33  	require.NoError(t, err)
    34  
    35  	// common execution for orm
    36  	mock.ExpectQuery("SELECT VERSION()").WillReturnRows(sqlmock.NewRows(
    37  		[]string{"VERSION()"}).AddRow("5.7.35-log"))
    38  
    39  	gdb, err := gorm.Open(mysql.New(mysql.Config{
    40  		Conn:                      db,
    41  		SkipInitializeWithVersion: false,
    42  	}), &gorm.Config{
    43  		SkipDefaultTransaction: true,
    44  	})
    45  	require.NoError(t, err)
    46  
    47  	return gdb, mock
    48  }
    49  
    50  func closeGormDB(t *testing.T, gdb *gorm.DB) {
    51  	db, err := gdb.DB()
    52  	require.NoError(t, err)
    53  	require.NoError(t, db.Close())
    54  }
    55  
    56  func TestNewEpochClient(t *testing.T) {
    57  	gdb, mock := mockGetDBConn(t)
    58  	defer closeGormDB(t, gdb)
    59  	defer mock.ExpectClose()
    60  
    61  	cli, err := NewEpochClient("fakeJob", gdb)
    62  	require.NoError(t, err)
    63  	defer cli.Close()
    64  }
    65  
    66  func TestGenEpoch(t *testing.T) {
    67  	gdb, mock := mockGetDBConn(t)
    68  	defer closeGormDB(t, gdb)
    69  	ctx, cancel := context.WithTimeout(context.TODO(), 1*time.Second)
    70  	defer cancel()
    71  
    72  	tm := time.Now()
    73  	createdAt := tm.Add(time.Duration(1))
    74  	updatedAt := tm.Add(time.Duration(1))
    75  
    76  	epochClient, err := NewEpochClient("fakeJob", gdb)
    77  	require.NoError(t, err)
    78  
    79  	// insert first record fail
    80  	mock.ExpectExec(".*").
    81  		WillReturnError(&gsql.MySQLError{Number: 1062, Message: "test error"})
    82  	epoch, err := epochClient.GenEpoch(ctx)
    83  	require.Error(t, err)
    84  	require.Equal(t, int64(0), epoch)
    85  	require.False(t, epochClient.isInitialized.Load())
    86  
    87  	// insert first record successful
    88  	mock.ExpectExec(regexp.QuoteMeta("INSERT INTO `logic_epoches` (`created_at`,`updated_at`,`job_id`,`epoch`)" +
    89  		" VALUES (?,?,?,?) ON DUPLICATE KEY UPDATE `seq_id`=`seq_id`")).
    90  		WillReturnResult(sqlmock.NewResult(1, 1))
    91  	mock.ExpectBegin()
    92  	mock.ExpectExec(regexp.QuoteMeta("UPDATE `logic_epoches` SET `epoch`=epoch + ?,`updated_at`=? WHERE job_id = ?")).
    93  		WillReturnResult(sqlmock.NewResult(1, 1))
    94  	mock.ExpectQuery(regexp.QuoteMeta("SELECT * FROM `logic_epoches` WHERE job_id = ? ORDER BY `logic_epoches`.`seq_id` LIMIT 1")).
    95  		WithArgs("fakeJob").
    96  		WillReturnRows(sqlmock.NewRows([]string{"seq_id", "created_at", "updated_at", "job_id", "epoch"}).
    97  			AddRow(1, createdAt, updatedAt, "fakeJob", 11))
    98  	mock.ExpectCommit()
    99  
   100  	epoch, err = epochClient.GenEpoch(ctx)
   101  	require.NoError(t, err)
   102  	require.Equal(t, int64(11), epoch)
   103  	require.True(t, epochClient.isInitialized.Load())
   104  
   105  	// update fail
   106  	mock.ExpectBegin()
   107  	mock.ExpectExec("UPDATE `logic_epoches` SET").WillReturnError(errors.New("gen epoch error"))
   108  	mock.ExpectRollback()
   109  	_, err = epochClient.GenEpoch(ctx)
   110  	require.Error(t, err)
   111  
   112  	// context cancel
   113  	ctx, cancel = context.WithTimeout(context.TODO(), 1*time.Second)
   114  	defer cancel()
   115  
   116  	err = failpoint.Enable("github.com/pingcap/tiflow/engine/pkg/orm/model/genEpochDelay", "sleep(2000)")
   117  	require.NoError(t, err)
   118  	ctx = failpoint.WithHook(ctx, func(ctx context.Context, fpname string) bool {
   119  		return ctx.Value(fpname) != nil
   120  	})
   121  	ctx2 := context.WithValue(ctx, "github.com/pingcap/tiflow/engine/pkg/orm/model/genEpochDelay", struct{}{})
   122  
   123  	_, err = epochClient.GenEpoch(ctx2)
   124  	require.Error(t, err)
   125  	require.Regexp(t, "context deadline exceed", err.Error())
   126  	failpoint.Disable("github.com/pingcap/tiflow/engine/pkg/orm/model/genEpochDelay")
   127  
   128  	mock.ExpectClose()
   129  }