package nodegroup

import (
	"context"
	"errors"
	"strings"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	mysqlDriver "github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dbmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/dbmocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func newStateMockDB(t *testing.T) (*gorm.DB, sqlmock.Sqlmock) {
	t.Helper()
	sqlDB, mockDB, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
	require.NoError(t, err)

	gdb, err := gorm.Open(mysql.New(mysql.Config{Conn: sqlDB, SkipInitializeWithVersion: true}), &gorm.Config{SkipDefaultTransaction: true})
	require.NoError(t, err)

	t.Cleanup(func() {
		mockDB.ExpectClose()
		require.NoError(t, sqlDB.Close())
	})

	return gdb, mockDB
}

func TestDataplaneNodeInstanceGroupStateRepo_Save(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name      string
		state     *entity.DataplaneNodeInstanceGroupState
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantID    bool
	}{
		{
			name:     "nil state",
			state:    nil,
			wantCode: codes.InvalidArgument,
		},
		{
			name: "generate id and save",
			state: &entity.DataplaneNodeInstanceGroupState{
				DataplaneNodeInstanceGroupState: schema.DataplaneNodeInstanceGroupState{
					DataplaneName:      "dp-1",
					NodeSpec:           "8C16G",
					PrimaryInstance:    "m5.2xlarge",
					State:              schema.NodeInstanceGroupStateNormal,
					StateEnteredAt:     now,
					FencingToken:       0,
					LastShortageReason: "shortage",
				},
			},
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(1, 1))
			},
			wantCode: codes.OK,
			wantID:   true,
		},
		{
			name: "duplicate key error",
			state: &entity.DataplaneNodeInstanceGroupState{
				DataplaneNodeInstanceGroupState: schema.DataplaneNodeInstanceGroupState{
					ID:              "state-1",
					DataplaneName:   "dp-1",
					NodeSpec:        "8C16G",
					PrimaryInstance: "m5.2xlarge",
					State:           schema.NodeInstanceGroupStateNormal,
				},
			},
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnError(gorm.ErrDuplicatedKey)
			},
			wantCode: codes.Internal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.Save(ctx, tt.state)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				if tt.wantID {
					assert.NotEmpty(t, tt.state.ID)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_Load(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name      string
		id        string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantNil   bool
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name: "not found",
			id:   "state-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("state-1", 1).WillReturnRows(sqlmock.NewRows([]string{"id"}))
			},
			wantCode: codes.NotFound,
		},
		{
			name: "query error",
			id:   "state-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("state-2", 1).WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name: "success",
			id:   "state-3",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-3", "dp-1", "8c16g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(1), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("state-3", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			state, err := repo.Load(ctx, tt.id)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				require.NotNil(t, state)
				assert.Equal(t, tt.id, state.ID)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, state)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_LoadByDataplaneAndSpec(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name            string
		dataplaneName   string
		nodeSpec        string
		primaryInstance string
		setupMock       func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode        codes.Code
		wantNil         bool
	}{
		{
			name:            "invalid args",
			dataplaneName:   "",
			nodeSpec:        "",
			primaryInstance: "",
			wantCode:        codes.InvalidArgument,
		},
		{
			name:            "not found",
			dataplaneName:   "dp-1",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"id"}))
			},
			wantCode: codes.OK,
			wantNil:  true,
		},
		{
			name:            "query error",
			dataplaneName:   "dp-1",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:            "success",
			dataplaneName:   "dp-1",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-1", "dp-1", "8c16g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(2), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			state, err := repo.LoadByDataplaneAndSpec(ctx, tt.dataplaneName, tt.nodeSpec, tt.primaryInstance)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				if tt.wantNil {
					assert.Nil(t, state)
				} else {
					require.NotNil(t, state)
					assert.Equal(t, tt.dataplaneName, state.DataplaneName)
					assert.Equal(t, strings.ToUpper(tt.nodeSpec), state.NodeSpec)
					assert.Equal(t, tt.primaryInstance, state.PrimaryInstance)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, state)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_LoadOrCreate(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name            string
		dataplaneName   string
		nodeSpec        string
		primaryInstance string
		setupMock       func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode        codes.Code
		wantNil         bool
	}{
		{
			name:            "invalid args",
			dataplaneName:   "",
			nodeSpec:        "",
			primaryInstance: "",
			wantCode:        codes.InvalidArgument,
		},
		{
			name:            "existing state",
			dataplaneName:   "dp-1",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-1", "dp-1", "8c16g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(1), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "8c16g", "m5.2xlarge", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
		},
		{
			name:            "create new state",
			dataplaneName:   "dp-1",
			nodeSpec:        "16C32G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "16c32g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"id"}))
				mockDB.ExpectExec("INSERT .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(1, 1))
			},
			wantCode: codes.OK,
		},
		{
			name:            "duplicate key on create",
			dataplaneName:   "dp-1",
			nodeSpec:        "32C64G",
			primaryInstance: "m5.2xlarge",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "32c64g", "m5.2xlarge", 1).WillReturnRows(sqlmock.NewRows([]string{"id"}))
				mockDB.ExpectExec("INSERT .*dataplane_node_instance_group_states.*").WillReturnError(&mysqlDriver.MySQLError{Number: 1062, Message: "Duplicate entry"})
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-dup", "dp-1", "32c64g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(0), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1", "32c64g", "m5.2xlarge", 1).WillReturnRows(rows)
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			state, err := repo.LoadOrCreate(ctx, tt.dataplaneName, tt.nodeSpec, tt.primaryInstance)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				require.NotNil(t, state)
				assert.Equal(t, tt.dataplaneName, state.DataplaneName)
				assert.Equal(t, strings.ToUpper(tt.nodeSpec), state.NodeSpec)
				assert.Equal(t, tt.primaryInstance, state.PrimaryInstance)
				assert.Equal(t, schema.NodeInstanceGroupStateNormal, state.State)
				assert.False(t, state.StateEnteredAt.IsZero())
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, state)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_ListByDataplane(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name          string
		dataplaneName string
		setupMock     func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode      codes.Code
		wantLen       int
	}{
		{
			name:          "empty dataplane",
			dataplaneName: "",
			wantCode:      codes.InvalidArgument,
		},
		{
			name:          "query error",
			dataplaneName: "dp-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:          "success",
			dataplaneName: "dp-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-1", "dp-1", "8c16g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(1), "", nil, now, now).
					AddRow("state-2", "dp-1", "16c32g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(2), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs("dp-1").WillReturnRows(rows)
			},
			wantCode: codes.OK,
			wantLen:  2,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			states, err := repo.ListByDataplane(ctx, tt.dataplaneName)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				assert.Len(t, states, tt.wantLen)
				for _, state := range states {
					assert.Equal(t, strings.ToUpper(state.NodeSpec), state.NodeSpec)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, states)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_ListByState(t *testing.T) {
	ctx := context.Background()
	now := time.Now()
	tests := []struct {
		name      string
		state     string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
		wantLen   int
	}{
		{
			name:  "query error",
			state: schema.NodeInstanceGroupStateNormal,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs(schema.NodeInstanceGroupStateNormal).WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name:  "success",
			state: schema.NodeInstanceGroupStateNormal,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().ReadDB(mock.Anything).Return(gdb)
				rows := sqlmock.NewRows([]string{
					"id",
					"dataplane_name",
					"node_spec",
					"primary_instance",
					"state",
					"state_entered_at",
					"fencing_token",
					"last_shortage_reason",
					"cooldown_until",
					"created_at",
					"updated_at",
				}).AddRow("state-1", "dp-1", "8c16g", "m5.2xlarge", schema.NodeInstanceGroupStateNormal, now, int64(1), "", nil, now, now)
				mockDB.ExpectQuery("SELECT .*dataplane_node_instance_group_states.*").WithArgs(schema.NodeInstanceGroupStateNormal).WillReturnRows(rows)
			},
			wantCode: codes.OK,
			wantLen:  1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			states, err := repo.ListByState(ctx, tt.state)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
				assert.Len(t, states, tt.wantLen)
				for _, state := range states {
					assert.Equal(t, strings.ToUpper(state.NodeSpec), state.NodeSpec)
				}
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
				assert.Nil(t, states)
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_Delete(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		name      string
		id        string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name: "delete error",
			id:   "state-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*deleted_at.*").WithArgs(sqlmock.AnyArg(), "state-1").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name: "success",
			id:   "state-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*deleted_at.*").WithArgs(sqlmock.AnyArg(), "state-2").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.Delete(ctx, tt.id)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_TransitionState(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		name         string
		id           string
		expected     string
		newState     string
		setupMock    func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode     codes.Code
		wantErrorNil bool
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:     "success",
			id:       "state-1",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
		{
			name:     "failed precondition",
			id:       "state-2",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantCode: codes.FailedPrecondition,
		},
		{
			name:     "update error",
			id:       "state-3",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.TransitionState(ctx, tt.id, tt.expected, tt.newState)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_SetCooldown(t *testing.T) {
	ctx := context.Background()
	until := time.Now().Add(time.Minute)
	tests := []struct {
		name      string
		id        string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name: "not found",
			id:   "state-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantCode: codes.NotFound,
		},
		{
			name: "update error",
			id:   "state-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name: "success",
			id:   "state-3",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.SetCooldown(ctx, tt.id, until)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_ClearCooldown(t *testing.T) {
	ctx := context.Background()
	tests := []struct {
		name      string
		id        string
		setupMock func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode  codes.Code
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name: "not found",
			id:   "state-1",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantCode: codes.NotFound,
		},
		{
			name: "update error",
			id:   "state-2",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
		{
			name: "success",
			id:   "state-3",
			setupMock: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setupMock != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setupMock(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.ClearCooldown(ctx, tt.id)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}

func TestDataplaneNodeInstanceGroupStateRepo_TransitionStateWithCooldown(t *testing.T) {
	ctx := context.Background()
	cooldown := time.Now().Add(time.Minute)
	tests := []struct {
		name     string
		id       string
		expected string
		newState string
		cooldown *time.Time
		setup    func(*gorm.DB, sqlmock.Sqlmock, *dbmocks.Provider)
		wantCode codes.Code
	}{
		{
			name:     "empty id",
			id:       "",
			wantCode: codes.InvalidArgument,
		},
		{
			name:     "success with cooldown",
			id:       "state-1",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			cooldown: &cooldown,
			setup: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
		{
			name:     "success with nil cooldown",
			id:       "state-2",
			expected: schema.NodeInstanceGroupStateShortageDetected,
			newState: schema.NodeInstanceGroupStateNormal,
			cooldown: nil,
			setup: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 1))
			},
			wantCode: codes.OK,
		},
		{
			name:     "failed precondition",
			id:       "state-3",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			cooldown: &cooldown,
			setup: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantCode: codes.FailedPrecondition,
		},
		{
			name:     "update error",
			id:       "state-4",
			expected: schema.NodeInstanceGroupStateNormal,
			newState: schema.NodeInstanceGroupStateShortageDetected,
			cooldown: &cooldown,
			setup: func(gdb *gorm.DB, mockDB sqlmock.Sqlmock, provider *dbmocks.Provider) {
				provider.EXPECT().WriteDB(mock.Anything).Return(gdb)
				mockDB.ExpectExec("UPDATE .*dataplane_node_instance_group_states.*").WillReturnError(errors.New("boom"))
			},
			wantCode: codes.Internal,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := dbmocks.NewProvider(t)
			var (
				gdb    *gorm.DB
				mockDB sqlmock.Sqlmock
			)
			if tt.setup != nil {
				gdb, mockDB = newStateMockDB(t)
				tt.setup(gdb, mockDB, provider)
			}

			repo := NewDataplaneNodeInstanceGroupStateRepo(provider)
			err := repo.TransitionStateWithCooldown(ctx, tt.id, tt.expected, tt.newState, tt.cooldown)
			if tt.wantCode == codes.OK {
				require.NoError(t, err)
			} else {
				require.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			}

			if mockDB != nil {
				require.NoError(t, mockDB.ExpectationsWereMet())
			}
		})
	}
}