@@ -2,7 +2,6 @@ package cli_test
22
33import (
44"bytes"
5- "context"
65"database/sql"
76"encoding/json"
87"io"
@@ -19,10 +18,7 @@ import (
1918"github.com/coder/coder/v2/cli/clitest"
2019"github.com/coder/coder/v2/coderd/coderdtest"
2120"github.com/coder/coder/v2/coderd/database"
22- "github.com/coder/coder/v2/coderd/database/dbauthz"
2321"github.com/coder/coder/v2/coderd/database/dbfake"
24- "github.com/coder/coder/v2/coderd/database/dbgen"
25- "github.com/coder/coder/v2/coderd/database/dbtime"
2622"github.com/coder/coder/v2/coderd/util/slice"
2723"github.com/coder/coder/v2/codersdk"
2824"github.com/coder/coder/v2/pty/ptytest"
@@ -43,76 +39,22 @@ func makeAITask(t *testing.T, db database.Store, orgID, adminID, ownerID uuid.UU
4339},
4440}).Do ()
4541
46- ws := database.WorkspaceTable {
42+ build := dbfake . WorkspaceBuild ( t , db , database.WorkspaceTable {
4743OrganizationID :orgID ,
4844OwnerID :ownerID ,
4945TemplateID :tv .Template .ID ,
50- }
51- build := dbfake .WorkspaceBuild (t ,db ,ws ).
46+ }).
5247Seed (database.WorkspaceBuild {
5348TemplateVersionID :tv .TemplateVersion .ID ,
5449Transition :transition ,
55- }).WithAgent ().Do ()
56- dbgen .WorkspaceBuildParameters (t ,db , []database.WorkspaceBuildParameter {
57- {
58- WorkspaceBuildID :build .Build .ID ,
59- Name :codersdk .AITaskPromptParameterName ,
60- Value :prompt ,
61- },
62- })
63- agents ,err := db .GetWorkspaceAgentsByWorkspaceAndBuildNumber (
64- dbauthz .AsSystemRestricted (context .Background ()),
65- database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams {
66- WorkspaceID :build .Workspace .ID ,
67- BuildNumber :build .Build .BuildNumber ,
68- },
69- )
70- require .NoError (t ,err )
71- require .NotEmpty (t ,agents )
72- agentID := agents [0 ].ID
73-
74- // Create a workspace app and set it as the sidebar app.
75- app := dbgen .WorkspaceApp (t ,db , database.WorkspaceApp {
76- AgentID :agentID ,
77- Slug :"task-sidebar" ,
78- DisplayName :"Task Sidebar" ,
79- External :false ,
80- })
81-
82- // Update build flags to reference the sidebar app and HasAITask=true.
83- err = db .UpdateWorkspaceBuildFlagsByID (
84- dbauthz .AsSystemRestricted (context .Background ()),
85- database.UpdateWorkspaceBuildFlagsByIDParams {
86- ID :build .Build .ID ,
87- HasAITask : sql.NullBool {Bool :true ,Valid :true },
88- HasExternalAgent : sql.NullBool {Bool :false ,Valid :false },
89- SidebarAppID : uuid.NullUUID {UUID :app .ID ,Valid :true },
90- UpdatedAt :build .Build .UpdatedAt ,
91- },
92- )
93- require .NoError (t ,err )
94-
95- // Create a task record in the tasks table for the new data model.
96- task := dbgen .Task (t ,db , database.TaskTable {
97- OrganizationID :orgID ,
98- OwnerID :ownerID ,
99- Name :build .Workspace .Name ,
100- WorkspaceID : uuid.NullUUID {UUID :build .Workspace .ID ,Valid :true },
101- TemplateVersionID :tv .TemplateVersion .ID ,
102- TemplateParameters : []byte ("{}" ),
103- Prompt :prompt ,
104- CreatedAt :dbtime .Now (),
105- })
106-
107- // Link the task to the workspace app.
108- dbgen .TaskWorkspaceApp (t ,db , database.TaskWorkspaceApp {
109- TaskID :task .ID ,
110- WorkspaceBuildNumber :build .Build .BuildNumber ,
111- WorkspaceAgentID : uuid.NullUUID {UUID :agentID ,Valid :true },
112- WorkspaceAppID : uuid.NullUUID {UUID :app .ID ,Valid :true },
113- })
114-
115- return task
50+ }).
51+ WithAgent ().
52+ WithTask (database.TaskTable {
53+ Prompt :prompt ,
54+ },nil ).
55+ Do ()
56+
57+ return build .Task
11658}
11759
11860func TestExpTaskList (t * testing.T ) {