Add unit test for helper function of sql execution store

common/persistence/sql/sql_execution_store_test.go

@@ -32,9 +32,11 @@ import (
  "github.com/stretchr/testify/require"
 
  "github.com/uber/cadence/common"
+ "github.com/uber/cadence/common/log/testlogger"
  "github.com/uber/cadence/common/persistence"
  "github.com/uber/cadence/common/persistence/serialization"
  "github.com/uber/cadence/common/persistence/sql/sqlplugin"
+ "github.com/uber/cadence/common/types"
 )
 
 func TestDeleteCurrentWorkflowExecution(t *testing.T) {
@@ -1883,3 +1885,73 @@ func TestDeleteWorkflowExecution(t *testing.T) {
  })
  }
 }
+
+func TestTxExecuteShardLocked(t *testing.T) {
+ tests := []struct {
+ name string
+ mockSetup func(*sqlplugin.MockDB, *sqlplugin.MockTx)
+ operation string
+ rangeID int64
+ fn func(sqlplugin.Tx) error
+ wantError error
+ }{
+ {
+ name: "Success",
+ mockSetup: func(mockDB *sqlplugin.MockDB, mockTx *sqlplugin.MockTx) {
+ mockDB.EXPECT().BeginTx(gomock.Any(), gomock.Any()).Return(mockTx, nil)
+ mockTx.EXPECT().ReadLockShards(gomock.Any(), gomock.Any()).Return(11, nil)
+ mockTx.EXPECT().Commit().Return(nil)
+ },
+ operation: "Insert",
+ rangeID: 11,
+ fn: func(sqlplugin.Tx) error { return nil },
+ wantError: nil,
+ },
+ {
+ name: "Error",
+ mockSetup: func(mockDB *sqlplugin.MockDB, mockTx *sqlplugin.MockTx) {
+ mockDB.EXPECT().BeginTx(gomock.Any(), gomock.Any()).Return(mockTx, nil)
+ mockTx.EXPECT().ReadLockShards(gomock.Any(), gomock.Any()).Return(11, nil)
+ mockTx.EXPECT().Rollback().Return(nil)
+ mockDB.EXPECT().IsNotFoundError(gomock.Any()).Return(false)
+ mockDB.EXPECT().IsTimeoutError(gomock.Any()).Return(false)
+ mockDB.EXPECT().IsThrottlingError(gomock.Any()).Return(false)
+ },
+ operation: "Insert",
+ rangeID: 11,
+ fn: func(sqlplugin.Tx) error { return errors.New("error") },
+ wantError: &types.InternalServiceError{Message: "Insert operation failed. Error: error"},
+ },
+ {
+ name: "Error - shard ownership lost",
+ mockSetup: func(mockDB *sqlplugin.MockDB, mockTx *sqlplugin.MockTx) {
+ mockDB.EXPECT().BeginTx(gomock.Any(), gomock.Any()).Return(mockTx, nil)
+ mockTx.EXPECT().ReadLockShards(gomock.Any(), gomock.Any()).Return(12, nil)
+ mockTx.EXPECT().Rollback().Return(nil)
+ },
+ operation: "Insert",
+ rangeID: 11,
+ fn: func(sqlplugin.Tx) error { return errors.New("error") },
+ wantError: &persistence.ShardOwnershipLostError{ShardID: 0, Msg: "Failed to lock shard. Previous range ID: 11; new range ID: 12"},
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ ctrl := gomock.NewController(t)
+ mockDB := sqlplugin.NewMockDB(ctrl)
+ mockTx := sqlplugin.NewMockTx(ctrl)
+ tt.mockSetup(mockDB, mockTx)
+
+ s := &sqlExecutionStore{
+ shardID: 0,
+ sqlStore: sqlStore{
+ db: mockDB,
+ logger: testlogger.New(t),
+ },
+ }
+
+ gotError := s.txExecuteShardLocked(context.Background(), 0, tt.operation, tt.rangeID, tt.fn)
+ assert.Equal(t, tt.wantError, gotError)
+ })
+ }
+}