Skip to content

Commit 2c6062f

Browse files
committed
Merge branch 'main' into rafrdz/fix-icon-removed
2 parents 7c1317e + ce93565 commit 2c6062f

File tree

8 files changed

+260
-26
lines changed

8 files changed

+260
-26
lines changed

coderd/database/dbauthz/dbauthz_test.go

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,11 @@ import (
1111
"testing"
1212
"time"
1313

14+
"github.com/brianvoe/gofakeit/v7"
1415
"github.com/google/uuid"
1516
"github.com/sqlc-dev/pqtype"
1617
"github.com/stretchr/testify/require"
18+
"go.uber.org/mock/gomock"
1719
"golang.org/x/xerrors"
1820

1921
"cdr.dev/slog"
@@ -22,6 +24,7 @@ import (
2224
"github.com/coder/coder/v2/coderd/database/db2sdk"
2325
"github.com/coder/coder/v2/coderd/database/dbauthz"
2426
"github.com/coder/coder/v2/coderd/database/dbgen"
27+
"github.com/coder/coder/v2/coderd/database/dbmock"
2528
"github.com/coder/coder/v2/coderd/database/dbtestutil"
2629
"github.com/coder/coder/v2/coderd/database/dbtime"
2730
"github.com/coder/coder/v2/coderd/notifications"
@@ -204,14 +207,15 @@ func defaultIPAddress() pqtype.Inet {
204207
}
205208

206209
func (s *MethodTestSuite) TestAPIKey() {
207-
s.Run("DeleteAPIKeyByID", s.Subtest(func(db database.Store, check *expects) {
208-
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
209-
key, _ := dbgen.APIKey(s.T(), db, database.APIKey{})
210+
s.Run("DeleteAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
211+
key := testutil.Fake(s.T(), faker, database.APIKey{})
212+
dbm.EXPECT().GetAPIKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes()
213+
dbm.EXPECT().DeleteAPIKeyByID(gomock.Any(), key.ID).Return(nil).AnyTimes()
210214
check.Args(key.ID).Asserts(key, policy.ActionDelete).Returns()
211215
}))
212-
s.Run("GetAPIKeyByID", s.Subtest(func(db database.Store, check *expects) {
213-
dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
214-
key, _ := dbgen.APIKey(s.T(), db, database.APIKey{})
216+
s.Run("GetAPIKeyByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
217+
key := testutil.Fake(s.T(), faker, database.APIKey{})
218+
dbm.EXPECT().GetAPIKeyByID(gomock.Any(), key.ID).Return(key, nil).AnyTimes()
215219
check.Args(key.ID).Asserts(key, policy.ActionRead).Returns(key)
216220
}))
217221
s.Run("GetAPIKeyByName", s.Subtest(func(db database.Store, check *expects) {
@@ -234,14 +238,12 @@ func (s *MethodTestSuite) TestAPIKey() {
234238
Asserts(a, policy.ActionRead, b, policy.ActionRead).
235239
Returns(slice.New(a, b))
236240
}))
237-
s.Run("GetAPIKeysByUserID", s.Subtest(func(db database.Store, check *expects) {
238-
u1 := dbgen.User(s.T(), db, database.User{})
239-
u2 := dbgen.User(s.T(), db, database.User{})
240-
241-
keyA, _ := dbgen.APIKey(s.T(), db, database.APIKey{UserID: u1.ID, LoginType: database.LoginTypeToken, TokenName: "key-a"})
242-
keyB, _ := dbgen.APIKey(s.T(), db, database.APIKey{UserID: u1.ID, LoginType: database.LoginTypeToken, TokenName: "key-b"})
243-
_, _ = dbgen.APIKey(s.T(), db, database.APIKey{UserID: u2.ID, LoginType: database.LoginTypeToken})
241+
s.Run("GetAPIKeysByUserID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
242+
u1 := testutil.Fake(s.T(), faker, database.User{})
243+
keyA := testutil.Fake(s.T(), faker, database.APIKey{UserID: u1.ID, LoginType: database.LoginTypeToken, TokenName: "key-a"})
244+
keyB := testutil.Fake(s.T(), faker, database.APIKey{UserID: u1.ID, LoginType: database.LoginTypeToken, TokenName: "key-b"})
244245

246+
dbm.EXPECT().GetAPIKeysByUserID(gomock.Any(), gomock.Any()).Return(slice.New(keyA, keyB), nil).AnyTimes()
245247
check.Args(database.GetAPIKeysByUserIDParams{LoginType: database.LoginTypeToken, UserID: u1.ID}).
246248
Asserts(keyA, policy.ActionRead, keyB, policy.ActionRead).
247249
Returns(slice.New(keyA, keyB))

coderd/database/dbauthz/setup_test.go

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"strings"
1111
"testing"
1212

13+
"github.com/brianvoe/gofakeit/v7"
1314
"github.com/google/go-cmp/cmp"
1415
"github.com/google/go-cmp/cmp/cmpopts"
1516
"github.com/google/uuid"
@@ -20,14 +21,14 @@ import (
2021
"golang.org/x/xerrors"
2122

2223
"cdr.dev/slog"
23-
"github.com/coder/coder/v2/coderd/rbac/policy"
2424

2525
"github.com/coder/coder/v2/coderd/coderdtest"
2626
"github.com/coder/coder/v2/coderd/database"
2727
"github.com/coder/coder/v2/coderd/database/dbauthz"
2828
"github.com/coder/coder/v2/coderd/database/dbmock"
2929
"github.com/coder/coder/v2/coderd/database/dbtestutil"
3030
"github.com/coder/coder/v2/coderd/rbac"
31+
"github.com/coder/coder/v2/coderd/rbac/policy"
3132
"github.com/coder/coder/v2/coderd/rbac/regosql"
3233
"github.com/coder/coder/v2/coderd/util/slice"
3334
)
@@ -105,19 +106,44 @@ func (s *MethodTestSuite) TearDownSuite() {
105106

106107
var testActorID = uuid.New()
107108

108-
// Subtest is a helper function that returns a function that can be passed to
109+
// Mocked runs a subtest with a mocked database. Removing the overhead of a real
110+
// postgres database resulting in much faster tests.
111+
func (s *MethodTestSuite) Mocked(testCaseF func(dmb *dbmock.MockStore, faker *gofakeit.Faker, check *expects)) func() {
112+
t := s.T()
113+
mDB := dbmock.NewMockStore(gomock.NewController(t))
114+
mDB.EXPECT().Wrappers().Return([]string{}).AnyTimes()
115+
116+
// Use a constant seed to prevent flakes from random data generation.
117+
faker := gofakeit.New(0)
118+
119+
// The usual Subtest assumes the test setup will use a real database to populate
120+
// with data. In this mocked case, we want to pass the underlying mocked database
121+
// to the test case instead.
122+
return s.SubtestWithDB(mDB, func(_ database.Store, check *expects) {
123+
testCaseF(mDB, faker, check)
124+
})
125+
}
126+
127+
// Subtest starts up a real postgres database for each test case.
128+
// Deprecated: Use 'Mocked' instead for much faster tests.
129+
func (s *MethodTestSuite) Subtest(testCaseF func(db database.Store, check *expects)) func() {
130+
t := s.T()
131+
db, _ := dbtestutil.NewDB(t)
132+
return s.SubtestWithDB(db, testCaseF)
133+
}
134+
135+
// SubtestWithDB is a helper function that returns a function that can be passed to
109136
// s.Run(). This function will run the test case for the method that is being
110137
// tested. The check parameter is used to assert the results of the method.
111138
// If the caller does not use the `check` parameter, the test will fail.
112-
func (s *MethodTestSuite) Subtest(testCaseF func(db database.Store, check *expects)) func() {
139+
func (s *MethodTestSuite) SubtestWithDB(db database.Store, testCaseF func(db database.Store, check *expects)) func() {
113140
return func() {
114141
t := s.T()
115142
testName := s.T().Name()
116143
names := strings.Split(testName, "/")
117144
methodName := names[len(names)-1]
118145
s.methodAccounting[methodName]++
119146

120-
db, _ := dbtestutil.NewDB(t)
121147
fakeAuthorizer := &coderdtest.FakeAuthorizer{}
122148
rec := &coderdtest.RecordingAuthorizer{
123149
Wrapped: fakeAuthorizer,

coderd/provisionerdserver/provisionerdserver.go

Lines changed: 58 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1925,12 +1925,16 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19251925
return xerrors.Errorf("update workspace build deadline: %w", err)
19261926
}
19271927

1928+
appIDs := make([]string, 0)
19281929
agentTimeouts := make(map[time.Duration]bool) // A set of agent timeouts.
19291930
// This could be a bulk insert to improve performance.
19301931
for _, protoResource := range jobType.WorkspaceBuild.Resources {
19311932
for _, protoAgent := range protoResource.Agents {
19321933
dur := time.Duration(protoAgent.GetConnectionTimeoutSeconds()) * time.Second
19331934
agentTimeouts[dur] = true
1935+
for _, app := range protoAgent.GetApps() {
1936+
appIDs = append(appIDs, app.GetId())
1937+
}
19341938
}
19351939

19361940
err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
@@ -1945,33 +1949,78 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19451949
}
19461950

19471951
var sidebarAppID uuid.NullUUID
1948-
hasAITask := len(jobType.WorkspaceBuild.AiTasks) == 1
1949-
if hasAITask {
1950-
task := jobType.WorkspaceBuild.AiTasks[0]
1951-
if task.SidebarApp == nil {
1952-
return xerrors.Errorf("update ai task: sidebar app is nil")
1952+
var hasAITask bool
1953+
var warnUnknownSidebarAppID bool
1954+
if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
1955+
hasAITask = true
1956+
task := tasks[0]
1957+
if task == nil || task.GetSidebarApp() == nil || len(task.GetSidebarApp().GetId()) == 0 {
1958+
return xerrors.Errorf("update ai task: sidebar app is nil or empty")
1959+
}
1960+
1961+
sidebarTaskID := task.GetSidebarApp().GetId()
1962+
if !slices.Contains(appIDs, sidebarTaskID) {
1963+
warnUnknownSidebarAppID = true
19531964
}
19541965

1955-
id, err := uuid.Parse(task.SidebarApp.Id)
1966+
id, err := uuid.Parse(task.GetSidebarApp().GetId())
19561967
if err != nil {
19571968
return xerrors.Errorf("parse sidebar app id: %w", err)
19581969
}
19591970

19601971
sidebarAppID = uuid.NullUUID{UUID: id, Valid: true}
19611972
}
19621973

1974+
if warnUnknownSidebarAppID {
1975+
// Ref: https://github.com/coder/coder/issues/18776
1976+
// This can happen for a number of reasons:
1977+
// 1. Misconfigured template
1978+
// 2. Count=0 on the agent due to stop transition, meaning the associated coder_app was not inserted.
1979+
// Failing the build at this point is not ideal, so log a warning instead.
1980+
s.Logger.Warn(ctx, "unknown ai_task_sidebar_app_id",
1981+
slog.F("ai_task_sidebar_app_id", sidebarAppID.UUID.String()),
1982+
slog.F("job_id", job.ID.String()),
1983+
slog.F("workspace_id", workspace.ID),
1984+
slog.F("workspace_build_id", workspaceBuild.ID),
1985+
slog.F("transition", string(workspaceBuild.Transition)),
1986+
)
1987+
// In order to surface this to the user, we will also insert a warning into the build logs.
1988+
if _, err := db.InsertProvisionerJobLogs(ctx, database.InsertProvisionerJobLogsParams{
1989+
JobID: jobID,
1990+
CreatedAt: []time.Time{now, now, now, now},
1991+
Source: []database.LogSource{database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon},
1992+
Level: []database.LogLevel{database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn},
1993+
Stage: []string{"Cleaning Up", "Cleaning Up", "Cleaning Up", "Cleaning Up"},
1994+
Output: []string{
1995+
fmt.Sprintf("Unknown ai_task_sidebar_app_id %q. This workspace will be unable to run AI tasks. This may be due to a template configuration issue, please check with the template author.", sidebarAppID.UUID.String()),
1996+
"Template author: double-check the following:",
1997+
" - You have associated the coder_ai_task with a valid coder_app in your template (ref: https://registry.terraform.io/providers/coder/coder/latest/docs/resources/ai_task).",
1998+
" - You have associated the coder_agent with at least one other compute resource. Agents with no other associated resources are not inserted into the database.",
1999+
},
2000+
}); err != nil {
2001+
s.Logger.Error(ctx, "insert provisioner job log for ai task sidebar app id warning",
2002+
slog.F("job_id", jobID),
2003+
slog.F("workspace_id", workspace.ID),
2004+
slog.F("workspace_build_id", workspaceBuild.ID),
2005+
slog.F("transition", string(workspaceBuild.Transition)),
2006+
)
2007+
}
2008+
// Important: reset hasAITask and sidebarAppID so that we don't run into a fk constraint violation.
2009+
hasAITask = false
2010+
sidebarAppID = uuid.NullUUID{}
2011+
}
2012+
19632013
// Regardless of whether there is an AI task or not, update the field to indicate one way or the other since it
19642014
// always defaults to nil. ONLY if has_ai_task=true MUST ai_task_sidebar_app_id be set.
1965-
err = db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
2015+
if err := db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
19662016
ID: workspaceBuild.ID,
19672017
HasAITask: sql.NullBool{
19682018
Bool: hasAITask,
19692019
Valid: true,
19702020
},
19712021
SidebarAppID: sidebarAppID,
19722022
UpdatedAt: now,
1973-
})
1974-
if err != nil {
2023+
}); err != nil {
19752024
return xerrors.Errorf("update workspace build ai tasks flag: %w", err)
19762025
}
19772026

coderd/provisionerdserver/provisionerdserver_test.go

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2794,6 +2794,22 @@ func TestCompleteJob(t *testing.T) {
27942794
},
27952795
expected: true,
27962796
},
2797+
// Checks regression for https://github.com/coder/coder/issues/18776
2798+
{
2799+
name: "non-existing app",
2800+
input: &proto.CompletedJob_WorkspaceBuild{
2801+
AiTasks: []*sdkproto.AITask{
2802+
{
2803+
Id: uuid.NewString(),
2804+
SidebarApp: &sdkproto.AITaskSidebarApp{
2805+
// Non-existing app ID would previously trigger a FK violation.
2806+
Id: uuid.NewString(),
2807+
},
2808+
},
2809+
},
2810+
},
2811+
expected: false,
2812+
},
27972813
} {
27982814
t.Run(tc.name, func(t *testing.T) {
27992815
t.Parallel()

go.mod

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -477,6 +477,7 @@ require (
477477
)
478478

479479
require (
480+
github.com/brianvoe/gofakeit/v7 v7.3.0
480481
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225
481482
github.com/coder/aisdk-go v0.0.9
482483
github.com/coder/preview v1.0.3

go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -830,6 +830,8 @@ github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl
830830
github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
831831
github.com/bramvdbogaerde/go-scp v1.5.0 h1:a9BinAjTfQh273eh7vd3qUgmBC+bx+3TRDtkZWmIpzM=
832832
github.com/bramvdbogaerde/go-scp v1.5.0/go.mod h1:on2aH5AxaFb2G0N5Vsdy6B0Ml7k9HuHSwfo1y0QzAbQ=
833+
github.com/brianvoe/gofakeit/v7 v7.3.0 h1:TWStf7/lLpAjKw+bqwzeORo9jvrxToWEwp9b1J2vApQ=
834+
github.com/brianvoe/gofakeit/v7 v7.3.0/go.mod h1:QXuPeBw164PJCzCUZVmgpgHJ3Llj49jSLVkKPMtxtxA=
833835
github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs=
834836
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
835837
github.com/bytecodealliance/wasmtime-go/v3 v3.0.2 h1:3uZCA/BLTIu+DqCfguByNMJa2HVHpXvjfy0Dy7g6fuA=

testutil/faker.go

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package testutil
2+
3+
import (
4+
"reflect"
5+
"testing"
6+
7+
"github.com/brianvoe/gofakeit/v7"
8+
"github.com/stretchr/testify/require"
9+
)
10+
11+
// Fake will populate any zero fields in the provided struct with fake data.
12+
// Non-zero fields will remain unchanged.
13+
// Usage:
14+
//
15+
// key := Fake(t, faker, database.APIKey{
16+
// TokenName: "keep-my-name",
17+
// })
18+
func Fake[T any](t *testing.T, faker *gofakeit.Faker, seed T) T {
19+
t.Helper()
20+
21+
var tmp T
22+
err := faker.Struct(&tmp)
23+
require.NoError(t, err, "failed to generate fake data for type %T", tmp)
24+
25+
mergeZero(&seed, tmp)
26+
return seed
27+
}
28+
29+
// mergeZero merges the fields of src into dst, but only if the field in dst is
30+
// currently the zero value.
31+
// Make sure `dst` is a pointer to a struct, otherwise the fields are not assignable.
32+
func mergeZero(dst any, src any) {
33+
srcv := reflect.ValueOf(src)
34+
if srcv.Kind() == reflect.Ptr {
35+
srcv = srcv.Elem()
36+
}
37+
remain := [][2]reflect.Value{
38+
{reflect.ValueOf(dst).Elem(), srcv},
39+
}
40+
41+
// Traverse the struct fields and set them only if they are currently zero.
42+
// This is a breadth-first traversal of the struct fields. Struct definitions
43+
// Should not be that deep, so we should not hit any stack overflow issues.
44+
for {
45+
if len(remain) == 0 {
46+
return
47+
}
48+
dv, sv := remain[0][0], remain[0][1]
49+
remain = remain[1:] //
50+
for i := 0; i < dv.NumField(); i++ {
51+
df := dv.Field(i)
52+
sf := sv.Field(i)
53+
if !df.CanSet() {
54+
continue
55+
}
56+
if df.IsZero() { // only write if currently zero
57+
df.Set(sf)
58+
continue
59+
}
60+
61+
if dv.Field(i).Kind() == reflect.Struct {
62+
// If the field is a struct, we need to traverse it as well.
63+
remain = append(remain, [2]reflect.Value{df, sf})
64+
}
65+
}
66+
}
67+
}

0 commit comments

Comments
 (0)