package store import ( "context" "errors" "os" "strings" "testing" "time" "github.com/google/uuid" ) func TestNormalizeIdentityBatchInput(t *testing.T) { id := uuid.NewString() input, err := normalizeIdentityBatchInput(IdentityBatchInput{ Action: " DISABLE ", IDs: []string{id, " " + id + " "}, }) if err != nil { t.Fatalf("normalize identity batch: %v", err) } if input.Action != "disable" || len(input.IDs) != 1 || input.IDs[0] != id { t.Fatalf("unexpected normalized batch: %+v", input) } for _, invalid := range []IdentityBatchInput{ {Action: "archive", IDs: []string{id}}, {Action: "enable", IDs: nil}, {Action: "delete", IDs: []string{"not-a-uuid"}}, } { if _, err := normalizeIdentityBatchInput(invalid); !errors.Is(err, ErrInvalidIdentityBatch) { t.Fatalf("invalid batch %+v error=%v", invalid, err) } } } func TestIdentityBatchOperationsAreAtomic(t *testing.T) { databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) if databaseURL == "" { t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity batch PostgreSQL integration tests") } ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) defer cancel() applyOIDCJITTestMigrations(t, ctx, databaseURL) db, err := Connect(ctx, databaseURL) if err != nil { t.Fatalf("connect identity batch test database: %v", err) } defer db.Close() suffix := strings.ReplaceAll(time.Now().UTC().Format("20060102150405.000000000"), ".", "") groupA, err := db.CreateUserGroup(ctx, UserGroupInput{GroupKey: "batch-a-" + suffix, Name: "Batch A", Source: "gateway", Status: "active"}) if err != nil { t.Fatalf("create group A: %v", err) } groupB, err := db.CreateUserGroup(ctx, UserGroupInput{GroupKey: "batch-b-" + suffix, Name: "Batch B", Source: "gateway", Status: "active"}) if err != nil { t.Fatalf("create group B: %v", err) } userA, err := db.CreateGatewayUser(ctx, GatewayUserInput{UserKey: "batch-user-a-" + suffix, Username: "batch-user-a-" + suffix, Source: "gateway", DefaultUserGroupID: groupA.ID, Status: "active"}) if err != nil { t.Fatalf("create user A: %v", err) } userB, err := db.CreateGatewayUser(ctx, GatewayUserInput{UserKey: "batch-user-b-" + suffix, Username: "batch-user-b-" + suffix, Source: "gateway", DefaultUserGroupID: groupB.ID, Status: "active"}) if err != nil { t.Fatalf("create user B: %v", err) } t.Cleanup(func() { _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_access_rules WHERE subject_id = ANY($1::uuid[])`, []string{groupA.ID, groupB.ID}) _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_users WHERE id = ANY($1::uuid[])`, []string{userA.ID, userB.ID}) _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_user_groups WHERE id = ANY($1::uuid[])`, []string{groupA.ID, groupB.ID}) }) if _, err := db.CreateAccessRule(ctx, AccessRuleInput{ SubjectType: "user_group", SubjectID: groupA.ID, ResourceType: "platform_model", ResourceID: uuid.NewString(), Effect: "deny", Status: "active", }); err != nil { t.Fatalf("create group access rule: %v", err) } result, err := db.BatchGatewayUsers(ctx, IdentityBatchInput{IDs: []string{userA.ID, userA.ID, userB.ID}, Action: "disable"}) if err != nil || result.AffectedCount != 2 || result.RequestedCount != 2 { t.Fatalf("disable users result=%+v err=%v", result, err) } assertIdentityStatuses(t, ctx, db, "gateway_users", []string{userA.ID, userB.ID}, "disabled") if _, err := db.BatchGatewayUsers(ctx, IdentityBatchInput{IDs: []string{userA.ID, userB.ID}, Action: "enable"}); err != nil { t.Fatalf("enable users: %v", err) } assertIdentityStatuses(t, ctx, db, "gateway_users", []string{userA.ID, userB.ID}, "active") if _, err := db.BatchUserGroups(ctx, IdentityBatchInput{IDs: []string{groupA.ID, uuid.NewString()}, Action: "disable"}); !errors.Is(err, ErrIdentityBatchTargetNotFound) { t.Fatalf("missing group batch error=%v", err) } assertIdentityStatuses(t, ctx, db, "gateway_user_groups", []string{groupA.ID}, "active") if _, err := db.BatchUserGroups(ctx, IdentityBatchInput{IDs: []string{groupA.ID, groupB.ID}, Action: "disable"}); err != nil { t.Fatalf("disable groups: %v", err) } assertIdentityStatuses(t, ctx, db, "gateway_user_groups", []string{groupA.ID, groupB.ID}, "disabled") if _, err := db.BatchUserGroups(ctx, IdentityBatchInput{IDs: []string{groupA.ID, groupB.ID}, Action: "delete"}); err != nil { t.Fatalf("delete groups: %v", err) } var groupCount, groupRuleCount, usersWithDeletedGroup int if err := db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_user_groups WHERE id = ANY($1::uuid[])`, []string{groupA.ID, groupB.ID}).Scan(&groupCount); err != nil { t.Fatalf("count deleted groups: %v", err) } if err := db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_access_rules WHERE subject_type = 'user_group' AND subject_id = ANY($1::uuid[])`, []string{groupA.ID, groupB.ID}).Scan(&groupRuleCount); err != nil { t.Fatalf("count deleted group rules: %v", err) } if err := db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_users WHERE id = ANY($1::uuid[]) AND default_user_group_id IS NOT NULL`, []string{userA.ID, userB.ID}).Scan(&usersWithDeletedGroup); err != nil { t.Fatalf("count stale default groups: %v", err) } if groupCount != 0 || groupRuleCount != 0 || usersWithDeletedGroup != 0 { t.Fatalf("group batch delete left data: groups=%d rules=%d userRefs=%d", groupCount, groupRuleCount, usersWithDeletedGroup) } if _, err := db.BatchGatewayUsers(ctx, IdentityBatchInput{IDs: []string{userA.ID, userB.ID}, Action: "delete"}); err != nil { t.Fatalf("delete users: %v", err) } var deletedUsers int if err := db.pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_users WHERE id = ANY($1::uuid[]) AND status = 'deleted' AND deleted_at IS NOT NULL`, []string{userA.ID, userB.ID}).Scan(&deletedUsers); err != nil { t.Fatalf("count deleted users: %v", err) } if deletedUsers != 2 { t.Fatalf("soft-deleted users=%d, want 2", deletedUsers) } } func assertIdentityStatuses(t *testing.T, ctx context.Context, db *Store, table string, ids []string, want string) { t.Helper() query := `SELECT COUNT(*) FROM ` + table + ` WHERE id = ANY($1::uuid[]) AND status = $2` var count int if err := db.pool.QueryRow(ctx, query, ids, want).Scan(&count); err != nil { t.Fatalf("read %s statuses: %v", table, err) } if count != len(ids) { t.Fatalf("%s status %s count=%d, want %d", table, want, count, len(ids)) } }