| 579 | } |
| 580 | |
| 581 | func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgent) database.WorkspaceAgent { |
| 582 | agt, err := db.InsertWorkspaceAgent(genCtx, database.InsertWorkspaceAgentParams{ |
| 583 | ID: takeFirst(orig.ID, uuid.New()), |
| 584 | ParentID: takeFirst(orig.ParentID, uuid.NullUUID{}), |
| 585 | CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()), |
| 586 | UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()), |
| 587 | Name: takeFirst(orig.Name, testutil.GetRandomName(t)), |
| 588 | ResourceID: takeFirst(orig.ResourceID, uuid.New()), |
| 589 | AuthToken: takeFirst(orig.AuthToken, uuid.New()), |
| 590 | AuthInstanceID: sql.NullString{ |
| 591 | String: takeFirst(orig.AuthInstanceID.String, testutil.GetRandomName(t)), |
| 592 | Valid: takeFirst(orig.AuthInstanceID.Valid, true), |
| 593 | }, |
| 594 | Architecture: takeFirst(orig.Architecture, "amd64"), |
| 595 | EnvironmentVariables: pqtype.NullRawMessage{ |
| 596 | RawMessage: takeFirstSlice(orig.EnvironmentVariables.RawMessage, []byte("{}")), |
| 597 | Valid: takeFirst(orig.EnvironmentVariables.Valid, false), |
| 598 | }, |
| 599 | OperatingSystem: takeFirst(orig.OperatingSystem, "linux"), |
| 600 | Directory: takeFirst(orig.Directory, ""), |
| 601 | InstanceMetadata: pqtype.NullRawMessage{ |
| 602 | RawMessage: takeFirstSlice(orig.ResourceMetadata.RawMessage, []byte("{}")), |
| 603 | Valid: takeFirst(orig.ResourceMetadata.Valid, false), |
| 604 | }, |
| 605 | ResourceMetadata: pqtype.NullRawMessage{ |
| 606 | RawMessage: takeFirstSlice(orig.ResourceMetadata.RawMessage, []byte("{}")), |
| 607 | Valid: takeFirst(orig.ResourceMetadata.Valid, false), |
| 608 | }, |
| 609 | ConnectionTimeoutSeconds: takeFirst(orig.ConnectionTimeoutSeconds, 3600), |
| 610 | TroubleshootingURL: takeFirst(orig.TroubleshootingURL, "https://example.com"), |
| 611 | MOTDFile: takeFirst(orig.MOTDFile, ""), |
| 612 | DisplayApps: append([]database.DisplayApp{}, orig.DisplayApps...), |
| 613 | DisplayOrder: takeFirst(orig.DisplayOrder, 1), |
| 614 | APIKeyScope: takeFirst(orig.APIKeyScope, database.AgentKeyScopeEnumAll), |
| 615 | }) |
| 616 | require.NoError(t, err, "insert workspace agent") |
| 617 | if orig.FirstConnectedAt.Valid || orig.LastConnectedAt.Valid || orig.DisconnectedAt.Valid || orig.LastConnectedReplicaID.Valid { |
| 618 | err = db.UpdateWorkspaceAgentConnectionByID(genCtx, database.UpdateWorkspaceAgentConnectionByIDParams{ |
| 619 | ID: agt.ID, |
| 620 | FirstConnectedAt: takeFirst(orig.FirstConnectedAt, agt.FirstConnectedAt), |
| 621 | LastConnectedAt: takeFirst(orig.LastConnectedAt, agt.LastConnectedAt), |
| 622 | DisconnectedAt: takeFirst(orig.DisconnectedAt, agt.DisconnectedAt), |
| 623 | LastConnectedReplicaID: takeFirst(orig.LastConnectedReplicaID, agt.LastConnectedReplicaID), |
| 624 | UpdatedAt: takeFirst(orig.UpdatedAt, agt.UpdatedAt), |
| 625 | }) |
| 626 | require.NoError(t, err, "update workspace agent first connected at") |
| 627 | } |
| 628 | |
| 629 | // If the lifecycle state is "ready", update the agent with the corresponding timestamps |
| 630 | if orig.LifecycleState == database.WorkspaceAgentLifecycleStateReady && orig.StartedAt.Valid && orig.ReadyAt.Valid { |
| 631 | err := db.UpdateWorkspaceAgentLifecycleStateByID(genCtx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ |
| 632 | ID: agt.ID, |
| 633 | LifecycleState: orig.LifecycleState, |
| 634 | StartedAt: orig.StartedAt, |
| 635 | ReadyAt: orig.ReadyAt, |
| 636 | }) |
| 637 | require.NoError(t, err, "update workspace agent lifecycle state") |
| 638 | } |