Skip to content

Commit 5d1bd5b

Browse files
mafredrimafredri
committed
add tests
1 parent 2036479 commit 5d1bd5b

File tree

2 files changed

+274
-1
lines changed

2 files changed

+274
-1
lines changed

‎cli/exp_tasklist_test.go

Lines changed: 273 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,273 @@
1+
package cli_test
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"database/sql"
7+
"encoding/json"
8+
"io"
9+
"testing"
10+
11+
"cdr.dev/slog"
12+
"cdr.dev/slog/sloggers/sloghuman"
13+
"github.com/google/uuid"
14+
"github.com/stretchr/testify/require"
15+
16+
"github.com/coder/coder/v2/cli/clitest"
17+
"github.com/coder/coder/v2/coderd/coderdtest"
18+
"github.com/coder/coder/v2/coderd/database"
19+
"github.com/coder/coder/v2/coderd/database/dbauthz"
20+
"github.com/coder/coder/v2/coderd/database/dbfake"
21+
"github.com/coder/coder/v2/coderd/database/dbgen"
22+
"github.com/coder/coder/v2/coderd/util/slice"
23+
"github.com/coder/coder/v2/codersdk"
24+
"github.com/coder/coder/v2/pty/ptytest"
25+
"github.com/coder/coder/v2/testutil"
26+
)
27+
28+
// makeAITask creates an AI-task workspace.
29+
func makeAITask(t *testing.T, db database.Store, orgID, adminID, ownerID uuid.UUID, transition database.WorkspaceTransition, prompt string) (workspace database.WorkspaceTable) {
30+
t.Helper()
31+
32+
tv := dbfake.TemplateVersion(t, db).
33+
Seed(database.TemplateVersion{
34+
OrganizationID: orgID,
35+
CreatedBy: adminID,
36+
HasAITask: sql.NullBool{
37+
Bool: true,
38+
Valid: true,
39+
},
40+
}).Do()
41+
42+
ws := database.WorkspaceTable{
43+
OrganizationID: orgID,
44+
OwnerID: ownerID,
45+
TemplateID: tv.Template.ID,
46+
}
47+
build := dbfake.WorkspaceBuild(t, db, ws).
48+
Seed(database.WorkspaceBuild{
49+
TemplateVersionID: tv.TemplateVersion.ID,
50+
Transition: transition,
51+
}).WithAgent().Do()
52+
dbgen.WorkspaceBuildParameters(t, db, []database.WorkspaceBuildParameter{
53+
{
54+
WorkspaceBuildID: build.Build.ID,
55+
Name: codersdk.AITaskPromptParameterName,
56+
Value: prompt,
57+
},
58+
})
59+
agents, err := db.GetWorkspaceAgentsByWorkspaceAndBuildNumber(
60+
dbauthz.AsSystemRestricted(context.Background()),
61+
database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
62+
WorkspaceID: build.Workspace.ID,
63+
BuildNumber: build.Build.BuildNumber,
64+
},
65+
)
66+
require.NoError(t, err)
67+
require.NotEmpty(t, agents)
68+
agentID := agents[0].ID
69+
70+
// Create a workspace app and set it as the sidebar app.
71+
app := dbgen.WorkspaceApp(t, db, database.WorkspaceApp{
72+
AgentID: agentID,
73+
Slug: "task-sidebar",
74+
DisplayName: "Task Sidebar",
75+
External: false,
76+
})
77+
78+
// Update build flags to reference the sidebar app and HasAITask=true.
79+
err = db.UpdateWorkspaceBuildFlagsByID(
80+
dbauthz.AsSystemRestricted(context.Background()),
81+
database.UpdateWorkspaceBuildFlagsByIDParams{
82+
ID: build.Build.ID,
83+
HasAITask: sql.NullBool{Bool: true, Valid: true},
84+
HasExternalAgent: sql.NullBool{Bool: false, Valid: false},
85+
SidebarAppID: uuid.NullUUID{UUID: app.ID, Valid: true},
86+
UpdatedAt: build.Build.UpdatedAt,
87+
},
88+
)
89+
require.NoError(t, err)
90+
91+
return build.Workspace
92+
}
93+
94+
func TestExpTaskList(t *testing.T) {
95+
t.Parallel()
96+
97+
t.Run("NoTasks_Table", func(t *testing.T) {
98+
t.Parallel()
99+
100+
// Quiet logger to reduce noise.
101+
quiet := slog.Make(sloghuman.Sink(io.Discard))
102+
client, _ := coderdtest.NewWithDatabase(t, &coderdtest.Options{Logger: &quiet})
103+
owner := coderdtest.CreateFirstUser(t, client)
104+
memberClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
105+
106+
inv, root := clitest.New(t, "exp", "task", "list")
107+
clitest.SetupConfig(t, memberClient, root)
108+
109+
pty := ptytest.New(t).Attach(inv)
110+
ctx := testutil.Context(t, testutil.WaitShort)
111+
112+
err := inv.WithContext(ctx).Run()
113+
require.NoError(t, err)
114+
115+
pty.ExpectMatch("No tasks found.")
116+
})
117+
118+
t.Run("Single_Table", func(t *testing.T) {
119+
t.Parallel()
120+
121+
// Quiet logger to reduce noise.
122+
quiet := slog.Make(sloghuman.Sink(io.Discard))
123+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{Logger: &quiet})
124+
owner := coderdtest.CreateFirstUser(t, client)
125+
memberClient, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
126+
127+
wantPrompt := "build me a web app"
128+
ws := makeAITask(t, db, owner.OrganizationID, owner.UserID, memberUser.ID, database.WorkspaceTransitionStart, wantPrompt)
129+
130+
inv, root := clitest.New(t, "exp", "task", "list", "--column", "id,name,status,initial prompt")
131+
clitest.SetupConfig(t, memberClient, root)
132+
133+
pty := ptytest.New(t).Attach(inv)
134+
ctx := testutil.Context(t, testutil.WaitShort)
135+
136+
err := inv.WithContext(ctx).Run()
137+
require.NoError(t, err)
138+
139+
// Validate the table includes the task and status.
140+
pty.ExpectMatch(ws.Name)
141+
pty.ExpectMatch("running")
142+
pty.ExpectMatch(wantPrompt)
143+
})
144+
145+
t.Run("StatusFilter_JSON", func(t *testing.T) {
146+
t.Parallel()
147+
148+
// Quiet logger to reduce noise.
149+
quiet := slog.Make(sloghuman.Sink(io.Discard))
150+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{Logger: &quiet})
151+
owner := coderdtest.CreateFirstUser(t, client)
152+
memberClient, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
153+
154+
// Create two AI tasks: one running, one stopped.
155+
running := makeAITask(t, db, owner.OrganizationID, owner.UserID, memberUser.ID, database.WorkspaceTransitionStart, "keep me running")
156+
stopped := makeAITask(t, db, owner.OrganizationID, owner.UserID, memberUser.ID, database.WorkspaceTransitionStop, "stop me please")
157+
158+
// Use JSON output to reliably validate filtering.
159+
inv, root := clitest.New(t, "exp", "task", "list", "--status=stopped", "--output=json")
160+
clitest.SetupConfig(t, memberClient, root)
161+
162+
ctx := testutil.Context(t, testutil.WaitShort)
163+
var stdout bytes.Buffer
164+
inv.Stdout = &stdout
165+
inv.Stderr = &stdout
166+
167+
err := inv.WithContext(ctx).Run()
168+
require.NoError(t, err)
169+
170+
var tasks []codersdk.Task
171+
require.NoError(t, json.Unmarshal(stdout.Bytes(), &tasks))
172+
173+
// Only the stopped task is returned.
174+
require.Len(t, tasks, 1, "expected one task after filtering")
175+
require.Equal(t, stopped.ID, tasks[0].ID)
176+
require.NotEqual(t, running.ID, tasks[0].ID)
177+
})
178+
179+
t.Run("UserFlag_Me_Table", func(t *testing.T) {
180+
t.Parallel()
181+
182+
quiet := slog.Make(sloghuman.Sink(io.Discard))
183+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{Logger: &quiet})
184+
owner := coderdtest.CreateFirstUser(t, client)
185+
_, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
186+
187+
_ = makeAITask(t, db, owner.OrganizationID, owner.UserID, memberUser.ID, database.WorkspaceTransitionStart, "other-task")
188+
ws := makeAITask(t, db, owner.OrganizationID, owner.UserID, owner.UserID, database.WorkspaceTransitionStart, "me-task")
189+
190+
inv, root := clitest.New(t, "exp", "task", "list", "--user", "me")
191+
clitest.SetupConfig(t, client, root)
192+
193+
pty := ptytest.New(t).Attach(inv)
194+
ctx := testutil.Context(t, testutil.WaitShort)
195+
196+
err := inv.WithContext(ctx).Run()
197+
require.NoError(t, err)
198+
199+
pty.ExpectMatch(ws.Name)
200+
})
201+
}
202+
203+
func TestExpTaskList_OwnerCanListOthers(t *testing.T) {
204+
t.Parallel()
205+
206+
// Quiet logger to reduce noise.
207+
quiet := slog.Make(sloghuman.Sink(io.Discard))
208+
ownerClient, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{Logger: &quiet})
209+
owner := coderdtest.CreateFirstUser(t, ownerClient)
210+
211+
// Create two additional members in the owner's organization.
212+
_, memberAUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
213+
_, memberBUser := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
214+
215+
// Seed an AI task for member A and B.
216+
_ = makeAITask(t, db, owner.OrganizationID, owner.UserID, memberAUser.ID, database.WorkspaceTransitionStart, "member-A-task")
217+
_ = makeAITask(t, db, owner.OrganizationID, owner.UserID, memberBUser.ID, database.WorkspaceTransitionStart, "member-B-task")
218+
219+
t.Run("OwnerListsSpecificUserWithUserFlag_JSON", func(t *testing.T) {
220+
t.Parallel()
221+
222+
// As the owner, list only member A tasks.
223+
inv, root := clitest.New(t, "exp", "task", "list", "--user", memberAUser.Username, "--output=json")
224+
clitest.SetupConfig(t, ownerClient, root)
225+
226+
var stdout bytes.Buffer
227+
inv.Stdout = &stdout
228+
229+
ctx := testutil.Context(t, testutil.WaitShort)
230+
231+
err := inv.WithContext(ctx).Run()
232+
require.NoError(t, err)
233+
234+
var tasks []codersdk.Task
235+
require.NoError(t, json.Unmarshal(stdout.Bytes(), &tasks))
236+
237+
// Expect at least one task and all tasks to belong to member A.
238+
require.NotEmpty(t, tasks, "expected at least one task for member A")
239+
for _, task := range tasks {
240+
require.Equal(t, memberAUser.ID, task.OwnerID, "expected only member A tasks")
241+
}
242+
})
243+
244+
t.Run("OwnerListsAllWithAllFlag_JSON", func(t *testing.T) {
245+
t.Parallel()
246+
247+
// As the owner, list all tasks to verify both member tasks are present.
248+
// Use JSON output to reliably validate filtering.
249+
inv, root := clitest.New(t, "exp", "task", "list", "--all", "--output=json")
250+
clitest.SetupConfig(t, ownerClient, root)
251+
252+
var stdout bytes.Buffer
253+
inv.Stdout = &stdout
254+
255+
ctx := testutil.Context(t, testutil.WaitShort)
256+
257+
err := inv.WithContext(ctx).Run()
258+
require.NoError(t, err)
259+
260+
var tasks []codersdk.Task
261+
require.NoError(t, json.Unmarshal(stdout.Bytes(), &tasks))
262+
263+
// Expect at least two tasks and ensure both owners (member A and member B) are represented.
264+
require.GreaterOrEqual(t, len(tasks), 2, "expected two or more tasks in --all listing")
265+
266+
// Use slice.Find for concise existence checks.
267+
_, foundA := slice.Find(tasks, func(t codersdk.Task) bool { return t.OwnerID == memberAUser.ID })
268+
_, foundB := slice.Find(tasks, func(t codersdk.Task) bool { return t.OwnerID == memberBUser.ID })
269+
270+
require.True(t, foundA, "expected at least one task for member A in --all listing")
271+
require.True(t, foundB, "expected at least one task for member B in --all listing")
272+
})
273+
}

‎coderd/aitasks.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,7 @@ func (api *API) tasksList(rw http.ResponseWriter, r *http.Request) {
280280
// Ensure that we only include AI task workspaces in the results.
281281
filter.HasAITask = sql.NullBool{Valid: true, Bool: true}
282282

283-
if filter.OwnerUsername == "me" || filter.OwnerUsername == "" {
283+
if filter.OwnerUsername == "me" {
284284
filter.OwnerID = apiKey.UserID
285285
filter.OwnerUsername = ""
286286
}

0 commit comments

Comments
 (0)