diff --git a/internal/analytics/collector_test.go b/internal/analytics/collector_test.go index a467cb23d..0778b97d2 100644 --- a/internal/analytics/collector_test.go +++ b/internal/analytics/collector_test.go @@ -156,6 +156,10 @@ func (m *mockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return nil } +func (m *mockConfigStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + return nil +} + func (m *mockConfigStore) UpdatePurchasePlanTx(ctx context.Context, _ pgx.Tx, plan *config.PurchasePlan) error { return nil } diff --git a/internal/api/mocks_test.go b/internal/api/mocks_test.go index 733cf9f9c..59694d2b4 100644 --- a/internal/api/mocks_test.go +++ b/internal/api/mocks_test.go @@ -122,6 +122,11 @@ func (m *MockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return args.Error(0) } +func (m *MockConfigStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + args := m.Called(ctx, planID) + return args.Error(0) +} + // UpdatePurchasePlanTx falls back to UpdatePurchasePlan when no // expectation is registered so tests that only assert on the un-tx // variant stay green — same pattern as SavePurchaseExecutionTx. diff --git a/internal/config/interfaces.go b/internal/config/interfaces.go index e0186f26a..d95049762 100644 --- a/internal/config/interfaces.go +++ b/internal/config/interfaces.go @@ -22,6 +22,10 @@ type StoreInterface interface { CreatePurchasePlan(ctx context.Context, plan *PurchasePlan) error GetPurchasePlan(ctx context.Context, planID string) (*PurchasePlan, error) UpdatePurchasePlan(ctx context.Context, plan *PurchasePlan) error + // IncrementPlanCurrentStep atomically advances the ramp schedule for planID + // inside a SELECT FOR UPDATE transaction, preventing the concurrent-write + // lost-update race described in issue #1071. + IncrementPlanCurrentStep(ctx context.Context, planID string) error // UpdatePurchasePlanTx is the tx-accepting variant of UpdatePurchasePlan. // Used from createPlannedPurchases' WithTx block so the per-row // SavePurchaseExecutionTx writes and the plan's next_execution_date diff --git a/internal/config/store_postgres.go b/internal/config/store_postgres.go index 447b031f1..4a73fd2a9 100644 --- a/internal/config/store_postgres.go +++ b/internal/config/store_postgres.go @@ -433,21 +433,21 @@ func (s *PostgresStore) CreatePurchasePlan(ctx context.Context, plan *PurchasePl return nil } -// GetPurchasePlan retrieves a purchase plan by ID -func (s *PostgresStore) GetPurchasePlan(ctx context.Context, planID string) (*PurchasePlan, error) { - query := ` - SELECT id, name, enabled, auto_purchase, notification_days_before, - services, ramp_schedule, created_at, updated_at, - next_execution_date, last_execution_date, last_notification_sent - FROM purchase_plans - WHERE id = $1 - ` - +const purchasePlanSelectCols = ` + SELECT id, name, enabled, auto_purchase, notification_days_before, + services, ramp_schedule, created_at, updated_at, + next_execution_date, last_execution_date, last_notification_sent + FROM purchase_plans` + +// scanPurchasePlanRow deserialises one purchase_plans row returned by QueryRow +// or the first row of a query with FOR UPDATE. Extracted so GetPurchasePlan +// and IncrementPlanCurrentStep share the same scan logic. +func scanPurchasePlanRow(row pgx.Row) (*PurchasePlan, error) { var plan PurchasePlan var servicesJSON, rampScheduleJSON []byte var nextExecDate, lastExecDate, lastNotifSent sql.NullTime - err := s.db.QueryRow(ctx, query, planID).Scan( + err := row.Scan( &plan.ID, &plan.Name, &plan.Enabled, @@ -461,24 +461,17 @@ func (s *PostgresStore) GetPurchasePlan(ctx context.Context, planID string) (*Pu &lastExecDate, &lastNotifSent, ) - if err != nil { - if err == pgx.ErrNoRows { - return nil, fmt.Errorf("%w: purchase plan %s", ErrNotFound, planID) - } - return nil, fmt.Errorf("failed to get purchase plan: %w", err) + return nil, err } - // Unmarshal JSONB fields if err := json.Unmarshal(servicesJSON, &plan.Services); err != nil { return nil, fmt.Errorf("failed to unmarshal services: %w", err) } - if err := json.Unmarshal(rampScheduleJSON, &plan.RampSchedule); err != nil { return nil, fmt.Errorf("failed to unmarshal ramp_schedule: %w", err) } - // Handle nullable timestamps if nextExecDate.Valid { plan.NextExecutionDate = &nextExecDate.Time } @@ -488,10 +481,61 @@ func (s *PostgresStore) GetPurchasePlan(ctx context.Context, planID string) (*Pu if lastNotifSent.Valid { plan.LastNotificationSent = &lastNotifSent.Time } - return &plan, nil } +// GetPurchasePlan retrieves a purchase plan by ID +func (s *PostgresStore) GetPurchasePlan(ctx context.Context, planID string) (*PurchasePlan, error) { + query := purchasePlanSelectCols + ` WHERE id = $1` + plan, err := scanPurchasePlanRow(s.db.QueryRow(ctx, query, planID)) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, fmt.Errorf("%w: purchase plan %s", ErrNotFound, planID) + } + return nil, fmt.Errorf("failed to get purchase plan: %w", err) + } + return plan, nil +} + +// IncrementPlanCurrentStep atomically advances the ramp schedule for planID +// inside a transaction. The row is locked with SELECT FOR UPDATE so concurrent +// callers (overlapping Lambda invocations, multi-tick cron) cannot both read +// the same CurrentStep value and both write CurrentStep+1, skipping a step. +// Returns nil when the plan no longer exists (deleted between execution and +// progress update) so the caller is not penalised for a race it cannot control. +func (s *PostgresStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + return s.WithTx(ctx, func(tx pgx.Tx) error { + row := tx.QueryRow(ctx, purchasePlanSelectCols+` WHERE id = $1 FOR UPDATE`, planID) + plan, err := scanPurchasePlanRow(row) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil + } + return fmt.Errorf("failed to lock purchase plan %s: %w", planID, err) + } + + if !plan.RampSchedule.IsComplete() { + plan.RampSchedule.CurrentStep++ + } + + if !plan.RampSchedule.IsComplete() { + nextDate := plan.RampSchedule.GetNextPurchaseDate() + plan.NextExecutionDate = &nextDate + } else { + plan.NextExecutionDate = nil + } + + now := time.Now() + plan.LastExecutionDate = &now + // Refresh updated_at on every increment. The plan was read from the DB + // with its previous UpdatedAt, so UpdatePurchasePlanTx's zero-value + // guard would otherwise persist a stale updated_at timestamp. + plan.UpdatedAt = now + + return s.UpdatePurchasePlanTx(ctx, tx, plan) + }) +} + // UpdatePurchasePlan updates an existing purchase plan. Delegates to // UpdatePurchasePlanTx inside a single-call WithTx — keeps the public // surface unchanged for callers that don't need to bundle this with diff --git a/internal/config/store_postgres_increment_step_test.go b/internal/config/store_postgres_increment_step_test.go new file mode 100644 index 000000000..4e3549a89 --- /dev/null +++ b/internal/config/store_postgres_increment_step_test.go @@ -0,0 +1,211 @@ +package config + +// store_postgres_increment_step_test.go -- pgxmock regression tests for +// IncrementPlanCurrentStep, the atomic ramp-step advance added for issue #1071. +// +// The bug: updatePlanProgress previously did a plain GetPurchasePlan -> +// CurrentStep++ -> UpdatePurchasePlan read-modify-write with no row lock, so +// two overlapping Lambda invocations could both read CurrentStep=N and both +// write N+1, skipping a ramp step (a lost update on a money path). +// +// These tests pin the fix's contract: +// - the read happens inside a transaction (Begin) and carries FOR UPDATE, +// - CurrentStep is advanced by exactly one and that value is persisted, +// - a plan deleted mid-race is tolerated (returns nil, no spurious error), +// - a completed ramp clears next_execution_date instead of advancing. +// +// pgxmock uses regexp query matching (see newMock), so an expectation whose +// pattern requires "FOR UPDATE" only matches if the production query actually +// emits it -- that is the atomicity guard. Dropping FOR UPDATE from the store +// makes TestPGXMock_IncrementPlanCurrentStep_LocksAndAdvances fail. + +import ( + "context" + "database/sql" + "encoding/json" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/pashagolub/pgxmock/v4" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// purchasePlanRowCols mirrors the SELECT column list in purchasePlanSelectCols. +var purchasePlanRowCols = []string{ + "id", "name", "enabled", "auto_purchase", "notification_days_before", + "services", "ramp_schedule", "created_at", "updated_at", + "next_execution_date", "last_execution_date", "last_notification_sent", +} + +// rampStepArg is a pgxmock.Argument that unmarshals the ramp_schedule JSONB +// passed to the UPDATE and asserts its CurrentStep, proving the advance landed +// in the persisted row rather than only in memory. +type rampStepArg struct{ want int } + +func (a rampStepArg) Match(v interface{}) bool { + b, ok := v.([]byte) + if !ok { + return false + } + var rs RampSchedule + if err := json.Unmarshal(b, &rs); err != nil { + return false + } + return rs.CurrentStep == a.want +} + +// nullTimeArg matches a *time.Time UPDATE argument by presence (nil vs set), +// used to assert next_execution_date is cleared on a completed ramp. +type nullTimeArg struct{ wantNil bool } + +func (a nullTimeArg) Match(v interface{}) bool { + tp, ok := v.(*time.Time) + if !ok { + return false + } + return (tp == nil) == a.wantNil +} + +// afterArg matches a time.Time UPDATE argument that is strictly after the given +// instant, used to assert updated_at is refreshed rather than persisted stale. +type afterArg struct{ notBefore time.Time } + +func (a afterArg) Match(v interface{}) bool { + ts, ok := v.(time.Time) + if !ok { + return false + } + return ts.After(a.notBefore) +} + +const purchasePlanUpdateArgs = 11 + +// incrementUpdateArgs builds the WithArgs matcher list for the UPDATE issued by +// IncrementPlanCurrentStep, asserting the persisted CurrentStep at $7, a +// refreshed updated_at at $8 (> staleUpdatedAt), and the next_execution_date +// presence at $9 while leaving the rest as AnyArg. +func incrementUpdateArgs(wantStep int, staleUpdatedAt time.Time, wantNextNil bool) []interface{} { + args := anyArgsCfg(purchasePlanUpdateArgs) + args[6] = rampStepArg{want: wantStep} // ramp_schedule = $7 + args[7] = afterArg{notBefore: staleUpdatedAt} // updated_at = $8 + args[8] = nullTimeArg{wantNil: wantNextNil} // next_execution_date = $9 + return args +} + +func TestPGXMock_IncrementPlanCurrentStep_LocksAndAdvances(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + ctx := context.Background() + + now := time.Now().Truncate(time.Second) + svcJSON, err := json.Marshal(map[string]ServiceConfig{}) + require.NoError(t, err) + ramp := RampSchedule{ + Type: "weekly", + PercentPerStep: 25, + StepIntervalDays: 7, + CurrentStep: 1, + TotalSteps: 4, + StartDate: now, + } + rampJSON, err := json.Marshal(ramp) + require.NoError(t, err) + + // Seed updated_at with an old timestamp so the test can assert the store + // refreshes it on increment rather than persisting the stale value. + stale := now.AddDate(0, 0, -30) + rows := pgxmock.NewRows(purchasePlanRowCols).AddRow( + "plan-123", "Ramp Plan", true, true, 3, + svcJSON, rampJSON, now, stale, + sql.NullTime{Valid: false}, sql.NullTime{Valid: false}, sql.NullTime{Valid: false}, + ) + + // Begin -> SELECT ... FOR UPDATE -> UPDATE (CurrentStep advanced 1 -> 2, + // updated_at refreshed, next_execution_date still set since ramp not + // complete) -> Commit. + mock.ExpectBegin() + mock.ExpectQuery(`SELECT[\s\S]*FROM purchase_plans[\s\S]*WHERE id = \$1 FOR UPDATE`). + WithArgs("plan-123"). + WillReturnRows(rows) + mock.ExpectExec(`UPDATE purchase_plans`). + WithArgs(incrementUpdateArgs(2, stale, false)...). + WillReturnResult(pgxmock.NewResult("UPDATE", 1)) + mock.ExpectCommit() + + err = store.IncrementPlanCurrentStep(ctx, "plan-123") + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestPGXMock_IncrementPlanCurrentStep_PlanDeletedMidRace(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + ctx := context.Background() + + mock.ExpectBegin() + mock.ExpectQuery(`SELECT[\s\S]*FROM purchase_plans[\s\S]*WHERE id = \$1 FOR UPDATE`). + WithArgs("gone"). + WillReturnError(pgx.ErrNoRows) + mock.ExpectCommit() + + // A plan deleted between execution and progress update must not error: the + // caller cannot control that race and should not be penalised for it. + err := store.IncrementPlanCurrentStep(ctx, "gone") + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestPGXMock_IncrementPlanCurrentStep_CompletedRampClearsNextDate(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + ctx := context.Background() + + now := time.Now().Truncate(time.Second) + svcJSON, err := json.Marshal(map[string]ServiceConfig{}) + require.NoError(t, err) + // Already at the last step: CurrentStep == TotalSteps means IsComplete, so + // the step is not advanced and next_execution_date is cleared. + ramp := RampSchedule{Type: "weekly", PercentPerStep: 25, StepIntervalDays: 7, CurrentStep: 4, TotalSteps: 4, StartDate: now} + rampJSON, err := json.Marshal(ramp) + require.NoError(t, err) + + stale := now.AddDate(0, 0, -30) + rows := pgxmock.NewRows(purchasePlanRowCols).AddRow( + "plan-done", "Done Plan", true, true, 3, + svcJSON, rampJSON, now, stale, + sql.NullTime{Valid: true, Time: now}, sql.NullTime{Valid: false}, sql.NullTime{Valid: false}, + ) + + mock.ExpectBegin() + mock.ExpectQuery(`SELECT[\s\S]*FROM purchase_plans[\s\S]*WHERE id = \$1 FOR UPDATE`). + WithArgs("plan-done"). + WillReturnRows(rows) + // Step stays at 4 (already complete), updated_at refreshed, and + // next_execution_date is cleared. + mock.ExpectExec(`UPDATE purchase_plans`). + WithArgs(incrementUpdateArgs(4, stale, true)...). + WillReturnResult(pgxmock.NewResult("UPDATE", 1)) + mock.ExpectCommit() + + err = store.IncrementPlanCurrentStep(ctx, "plan-done") + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} + +func TestPGXMock_IncrementPlanCurrentStep_LockErrorRollsBack(t *testing.T) { + mock := newMock(t) + store := storeWith(mock) + ctx := context.Background() + + mock.ExpectBegin() + mock.ExpectQuery(`SELECT[\s\S]*FROM purchase_plans[\s\S]*WHERE id = \$1 FOR UPDATE`). + WithArgs("plan-err"). + WillReturnError(assert.AnError) + mock.ExpectRollback() + + err := store.IncrementPlanCurrentStep(ctx, "plan-err") + require.Error(t, err) + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/internal/mocks/stores.go b/internal/mocks/stores.go index aedb23d6c..3a883a35c 100644 --- a/internal/mocks/stores.go +++ b/internal/mocks/stores.go @@ -75,9 +75,15 @@ func (m *MockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return args.Error(0) } +// IncrementPlanCurrentStep mocks the atomic step-advance operation. +func (m *MockConfigStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + args := m.Called(ctx, planID) + return args.Error(0) +} + // UpdatePurchasePlanTx mocks the UpdatePurchasePlanTx operation. Falls // back to UpdatePurchasePlan when no expectation is registered so tests -// that don't care about the Tx variant stay green — same pattern as +// that don't care about the Tx variant stay green -- same pattern as // SavePurchaseExecutionTx below. func (m *MockConfigStore) UpdatePurchasePlanTx(ctx context.Context, tx pgx.Tx, plan *config.PurchasePlan) error { if !isExpected(&m.Mock, "UpdatePurchasePlanTx") { diff --git a/internal/purchase/approvals_test.go b/internal/purchase/approvals_test.go index 473a55eaf..4189845d8 100644 --- a/internal/purchase/approvals_test.go +++ b/internal/purchase/approvals_test.go @@ -45,8 +45,8 @@ func stubExecuteChain(t *testing.T, store *MockConfigStore, sender *MockEmailSen sender.On("SendPurchaseConfirmation", mock.Anything, mock.Anything).Return(nil) // finalizeExecution writes status=completed via SavePurchaseExecution. store.On("SavePurchaseExecution", mock.Anything, mock.AnythingOfType("*config.PurchaseExecution")).Return(nil) - // updatePlanProgress fetches+updates the plan. - store.On("UpdatePurchasePlan", mock.Anything, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + // updatePlanProgress calls IncrementPlanCurrentStep (atomic, issue #1071). + store.On("IncrementPlanCurrentStep", mock.Anything, planID).Return(nil) } func TestManager_ApproveExecution_Success(t *testing.T) { diff --git a/internal/purchase/coverage_extra_test.go b/internal/purchase/coverage_extra_test.go index ee88ca376..69dfbccfd 100644 --- a/internal/purchase/coverage_extra_test.go +++ b/internal/purchase/coverage_extra_test.go @@ -212,7 +212,7 @@ func TestHandleExecutePurchase_ApprovedStatus(t *testing.T) { mockStore.On("GetPurchasePlan", ctx, "plan-approved").Return(plan, nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")).Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-approved").Return(nil) mockSTS.On("GetCallerIdentity", ctx, mock.Anything).Return(nil, errors.New("sts error")) manager := &Manager{ @@ -368,7 +368,7 @@ func TestProcessMessage_ApproveHappyPath(t *testing.T) { mockStore.On("GetPurchasePlan", ctx, planID).Return(plan, nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.Anything).Return(nil) mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")).Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, planID).Return(nil) manager := &Manager{ config: mockStore, diff --git a/internal/purchase/execution.go b/internal/purchase/execution.go index 8ab258a6d..b1c159ace 100644 --- a/internal/purchase/execution.go +++ b/internal/purchase/execution.go @@ -1028,38 +1028,17 @@ func mapSavingsPlansSlug(service string) (common.ServiceType, bool) { } // updatePlanProgress advances the ramp schedule after a purchase. -// Direct-execute purchases (Opportunities flow) have no plan to advance — +// Direct-execute purchases (Opportunities flow) have no plan to advance -- // PlanID is empty and the Postgres UUID column would reject the query // with SQLSTATE 22P02, so short-circuit cleanly before hitting the store. +// The actual increment is delegated to IncrementPlanCurrentStep which holds +// a SELECT FOR UPDATE lock for the duration, preventing the lost-update race +// when two Lambda invocations overlap on the same plan (issue #1071). func (m *Manager) updatePlanProgress(ctx context.Context, planID string) error { if planID == "" { return nil } - plan, err := m.config.GetPurchasePlan(ctx, planID) - if err != nil { - return err - } - if plan == nil { - return nil - } - - // Advance ramp schedule only if not complete - if !plan.RampSchedule.IsComplete() { - plan.RampSchedule.CurrentStep++ - } - - // Calculate next execution date - if !plan.RampSchedule.IsComplete() { - nextDate := plan.RampSchedule.GetNextPurchaseDate() - plan.NextExecutionDate = &nextDate - } else { - plan.NextExecutionDate = nil - } - - now := time.Now() - plan.LastExecutionDate = &now - - return m.config.UpdatePurchasePlan(ctx, plan) + return m.config.IncrementPlanCurrentStep(ctx, planID) } // getAWSAccountID retrieves the current AWS account ID using STS. diff --git a/internal/purchase/execution_test.go b/internal/purchase/execution_test.go index 310daf894..6eb8ad387 100644 --- a/internal/purchase/execution_test.go +++ b/internal/purchase/execution_test.go @@ -320,22 +320,7 @@ func TestManager_UpdatePlanProgress(t *testing.T) { mockStore := new(MockConfigStore) mockEmail := new(MockEmailSender) - startDate := time.Now() - plan := &config.PurchasePlan{ - ID: "plan-123", - Name: "Test Plan", - RampSchedule: config.RampSchedule{ - Type: "weekly", - PercentPerStep: 25, - StepIntervalDays: 7, - CurrentStep: 0, - TotalSteps: 4, - StartDate: startDate, - }, - } - - mockStore.On("GetPurchasePlan", ctx, "plan-123").Return(plan, nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-123").Return(nil) manager := &Manager{ config: mockStore, @@ -354,7 +339,8 @@ func TestManager_UpdatePlanProgress_PlanNotFound(t *testing.T) { mockStore := new(MockConfigStore) mockEmail := new(MockEmailSender) - mockStore.On("GetPurchasePlan", ctx, "nonexistent").Return(nil, nil) + // IncrementPlanCurrentStep returns nil when the plan no longer exists. + mockStore.On("IncrementPlanCurrentStep", ctx, "nonexistent").Return(nil) manager := &Manager{ config: mockStore, @@ -373,7 +359,7 @@ func TestManager_UpdatePlanProgress_GetError(t *testing.T) { mockStore := new(MockConfigStore) mockEmail := new(MockEmailSender) - mockStore.On("GetPurchasePlan", ctx, "plan-123").Return(nil, errors.New("database error")) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-123").Return(errors.New("database error")) manager := &Manager{ config: mockStore, @@ -392,22 +378,9 @@ func TestManager_UpdatePlanProgress_CompleteRamp(t *testing.T) { mockStore := new(MockConfigStore) mockEmail := new(MockEmailSender) - startDate := time.Now() - plan := &config.PurchasePlan{ - ID: "plan-123", - Name: "Test Plan", - RampSchedule: config.RampSchedule{ - Type: "weekly", - PercentPerStep: 25, - StepIntervalDays: 7, - CurrentStep: 3, // Last step (0-indexed, so 3 is the 4th step) - TotalSteps: 4, - StartDate: startDate, - }, - } - - mockStore.On("GetPurchasePlan", ctx, "plan-123").Return(plan, nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + // The step-advance and completion logic now live inside IncrementPlanCurrentStep + // in the store, tested at the store layer with pgxmock. + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-123").Return(nil) manager := &Manager{ config: mockStore, diff --git a/internal/purchase/manager_test.go b/internal/purchase/manager_test.go index 50466f17f..c1386d8c6 100644 --- a/internal/purchase/manager_test.go +++ b/internal/purchase/manager_test.go @@ -212,11 +212,11 @@ func TestManager_ProcessScheduledPurchases_DuePurchase(t *testing.T) { mockStore.On("GetPendingExecutions", ctx).Return(executions, nil) mockStore.On("TransitionExecutionStatus", ctx, "exec-123", []string{"approved", "pending", "notified"}, "running").Return(&claimedExec, nil) - mockStore.On("GetPurchasePlan", ctx, "plan-456").Return(plan, nil).Twice() + mockStore.On("GetPurchasePlan", ctx, "plan-456").Return(plan, nil).Once() mockStore.On("SavePurchaseHistory", ctx, mock.AnythingOfType("*config.PurchaseHistoryRecord")).Return(nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")).Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-456").Return(nil) mockSTS.On("GetCallerIdentity", ctx, mock.AnythingOfType("*sts.GetCallerIdentityInput")).Return(&sts.GetCallerIdentityOutput{ Account: aws.String("123456789012"), }, nil) @@ -461,14 +461,14 @@ func TestManager_RecoverStrandedApprovals_AWSOnlyRedrives(t *testing.T) { // CAS claim: approved -> running. The re-drive proceeds only after winning this. mockStore.On("TransitionExecutionStatus", ctx, "exec-aws-stranded", []string{"approved"}, "running"). Return(&runningRow, nil) - mockStore.On("GetPurchasePlan", ctx, "plan-aws-456").Return(plan, nil).Twice() + mockStore.On("GetPurchasePlan", ctx, "plan-aws-456").Return(plan, nil).Once() mockStore.On("SavePurchaseHistory", ctx, mock.AnythingOfType("*config.PurchaseHistoryRecord")).Return(nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) var saved *config.PurchaseExecution mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")). Run(func(args mock.Arguments) { saved = args.Get(1).(*config.PurchaseExecution) }). Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-aws-456").Return(nil) mockSTS.On("GetCallerIdentity", ctx, mock.AnythingOfType("*sts.GetCallerIdentityInput")).Return(&sts.GetCallerIdentityOutput{ Account: aws.String("123456789012"), }, nil) @@ -549,14 +549,14 @@ func TestManager_RecoverStrandedApprovals_AzureReservationRedrives(t *testing.T) // CAS claim: approved -> running before re-drive. mockStore.On("TransitionExecutionStatus", ctx, "exec-azure-res-stranded", []string{"approved"}, "running"). Return(&runningRow, nil) - mockStore.On("GetPurchasePlan", ctx, "plan-azure-res").Return(plan, nil).Twice() + mockStore.On("GetPurchasePlan", ctx, "plan-azure-res").Return(plan, nil).Once() mockStore.On("SavePurchaseHistory", ctx, mock.AnythingOfType("*config.PurchaseHistoryRecord")).Return(nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) var saved *config.PurchaseExecution mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")). Run(func(args mock.Arguments) { saved = args.Get(1).(*config.PurchaseExecution) }). Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-azure-res").Return(nil) mockFactory.On("CreateAndValidateProvider", mock.Anything, "azure", mock.Anything).Return(mockProvider, nil) mockProvider.On("GetServiceClient", mock.Anything, common.ServiceCompute, "eastus").Return(mockServiceClient, nil) @@ -630,14 +630,14 @@ func TestManager_RecoverStrandedApprovals_GCPRedrives(t *testing.T) { // CAS claim: approved -> running before re-drive. mockStore.On("TransitionExecutionStatus", ctx, "exec-gcp-stranded", []string{"approved"}, "running"). Return(&runningRow, nil) - mockStore.On("GetPurchasePlan", ctx, "plan-gcp").Return(plan, nil).Twice() + mockStore.On("GetPurchasePlan", ctx, "plan-gcp").Return(plan, nil).Once() mockStore.On("SavePurchaseHistory", ctx, mock.AnythingOfType("*config.PurchaseHistoryRecord")).Return(nil) mockEmail.On("SendPurchaseConfirmation", ctx, mock.AnythingOfType("email.NotificationData")).Return(nil) var saved *config.PurchaseExecution mockStore.On("SavePurchaseExecution", ctx, mock.AnythingOfType("*config.PurchaseExecution")). Run(func(args mock.Arguments) { saved = args.Get(1).(*config.PurchaseExecution) }). Return(nil) - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil) + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-gcp").Return(nil) mockFactory.On("CreateAndValidateProvider", mock.Anything, "gcp", mock.Anything).Return(mockProvider, nil) mockProvider.On("GetServiceClient", mock.Anything, common.ServiceCompute, "us-central1").Return(mockServiceClient, nil) diff --git a/internal/purchase/mocks_test.go b/internal/purchase/mocks_test.go index d15c6897f..0331c7e9c 100644 --- a/internal/purchase/mocks_test.go +++ b/internal/purchase/mocks_test.go @@ -215,6 +215,11 @@ func (m *MockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return args.Error(0) } +func (m *MockConfigStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + args := m.Called(ctx, planID) + return args.Error(0) +} + // UpdatePurchasePlanTx falls back to UpdatePurchasePlan when no // expectation is registered. Mirrors the pattern used by // SavePurchaseExecutionTx and the api-package mock. diff --git a/internal/purchase/money_path_regression_test.go b/internal/purchase/money_path_regression_test.go index 047db4b28..7b38b95f4 100644 --- a/internal/purchase/money_path_regression_test.go +++ b/internal/purchase/money_path_regression_test.go @@ -318,7 +318,7 @@ func TestMultiAccountPartialSuccessIsAcked(t *testing.T) { // updatePlanProgress only runs on a fully-clean run (execErr == nil); a // partial multi-account run skips it. Marked Maybe so a (non-deterministic) // all-success ordering of the two account goroutines doesn't fail the mock. - mockStore.On("UpdatePurchasePlan", ctx, mock.AnythingOfType("*config.PurchasePlan")).Return(nil).Maybe() + mockStore.On("IncrementPlanCurrentStep", ctx, "plan-x").Return(nil).Maybe() var savedRoot *config.PurchaseExecution mockStore.SavePurchaseExecutionFn = func(_ context.Context, e *config.PurchaseExecution) error { diff --git a/internal/scheduler/scheduler_test.go b/internal/scheduler/scheduler_test.go index 1bfd3403d..074ceaeba 100644 --- a/internal/scheduler/scheduler_test.go +++ b/internal/scheduler/scheduler_test.go @@ -91,6 +91,11 @@ func (m *MockConfigStore) UpdatePurchasePlan(ctx context.Context, plan *config.P return args.Error(0) } +func (m *MockConfigStore) IncrementPlanCurrentStep(ctx context.Context, planID string) error { + args := m.Called(ctx, planID) + return args.Error(0) +} + // UpdatePurchasePlanTx falls back to UpdatePurchasePlan when no // expectation is registered. When an expectation is registered, the // transaction is forwarded to m.Called so test expectations can match diff --git a/internal/server/test_helpers_test.go b/internal/server/test_helpers_test.go index 2725d9f73..990cb131c 100644 --- a/internal/server/test_helpers_test.go +++ b/internal/server/test_helpers_test.go @@ -47,6 +47,10 @@ func (m *mockConfigStoreForHealth) UpdatePurchasePlan(ctx context.Context, plan return nil } +func (m *mockConfigStoreForHealth) IncrementPlanCurrentStep(_ context.Context, _ string) error { + return nil +} + func (m *mockConfigStoreForHealth) UpdatePurchasePlanTx(_ context.Context, _ pgx.Tx, _ *config.PurchasePlan) error { return nil }