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 }