Skip to content

Commit 9a5f639

Browse files
mafredrimafredri
committed
add tests
1 parent 2036479 commit 9a5f639

File tree

2 files changed

+279
-1
lines changed

2 files changed

+279
-1
lines changed

‎cli/exp_tasklist_test.go

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

‎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)