| 457 | } |
| 458 | |
| 459 | func GetRunningPrebuilds( |
| 460 | ctx context.Context, |
| 461 | t *testing.T, |
| 462 | db database.Store, |
| 463 | desiredPrebuilds int, |
| 464 | ) []database.GetRunningPrebuiltWorkspacesRow { |
| 465 | t.Helper() |
| 466 | |
| 467 | var runningPrebuilds []database.GetRunningPrebuiltWorkspacesRow |
| 468 | testutil.Eventually(ctx, t, func(context.Context) bool { |
| 469 | prebuiltWorkspaces, err := db.GetRunningPrebuiltWorkspaces(ctx) |
| 470 | assert.NoError(t, err, "failed to get running prebuilds") |
| 471 | |
| 472 | for _, prebuild := range prebuiltWorkspaces { |
| 473 | runningPrebuilds = append(runningPrebuilds, prebuild) |
| 474 | |
| 475 | agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, prebuild.ID) |
| 476 | assert.NoError(t, err, "failed to get agents") |
| 477 | |
| 478 | // Manually mark all agents as ready since tests don't have real agent processes |
| 479 | // that would normally report their lifecycle state. Prebuilt workspaces are only |
| 480 | // eligible for claiming when their agents reach the "ready" state. |
| 481 | for _, agent := range agents { |
| 482 | err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ |
| 483 | ID: agent.ID, |
| 484 | LifecycleState: database.WorkspaceAgentLifecycleStateReady, |
| 485 | StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, |
| 486 | ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, |
| 487 | }) |
| 488 | assert.NoError(t, err, "failed to update agent") |
| 489 | } |
| 490 | } |
| 491 | |
| 492 | t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), desiredPrebuilds) |
| 493 | return len(runningPrebuilds) == desiredPrebuilds |
| 494 | }, testutil.IntervalSlow, "found %d running prebuilds, expected %d", len(runningPrebuilds), desiredPrebuilds) |
| 495 | |
| 496 | return runningPrebuilds |
| 497 | } |
| 498 | |
| 499 | func MustRunReconciliationLoopForPreset( |
| 500 | ctx context.Context, |