Skip to content

Commit 2bc1070

Browse files
ssncferreirassncferreira
committed
Merge remote-tracking branch 'origin/main' into ssncferreira/fix-prebuilds-lifecycle-params
2 parents 4c847c9 + ce93565 commit 2bc1070

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
@@ -1941,12 +1941,16 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19411941
return xerrors.Errorf("update workspace build deadline: %w", err)
19421942
}
19431943

1944+
appIDs := make([]string, 0)
19441945
agentTimeouts := make(map[time.Duration]bool) // A set of agent timeouts.
19451946
// This could be a bulk insert to improve performance.
19461947
for _, protoResource := range jobType.WorkspaceBuild.Resources {
19471948
for _, protoAgent := range protoResource.Agents {
19481949
dur := time.Duration(protoAgent.GetConnectionTimeoutSeconds()) * time.Second
19491950
agentTimeouts[dur] = true
1951+
for _, app := range protoAgent.GetApps() {
1952+
appIDs = append(appIDs, app.GetId())
1953+
}
19501954
}
19511955

19521956
err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
@@ -1961,33 +1965,78 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19611965
}
19621966

19631967
var sidebarAppID uuid.NullUUID
1964-
hasAITask := len(jobType.WorkspaceBuild.AiTasks) == 1
1965-
if hasAITask {
1966-
task := jobType.WorkspaceBuild.AiTasks[0]
1967-
if task.SidebarApp == nil {
1968-
return xerrors.Errorf("update ai task: sidebar app is nil")
1968+
var hasAITask bool
1969+
var warnUnknownSidebarAppID bool
1970+
if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
1971+
hasAITask = true
1972+
task := tasks[0]
1973+
if task == nil || task.GetSidebarApp() == nil || len(task.GetSidebarApp().GetId()) == 0 {
1974+
return xerrors.Errorf("update ai task: sidebar app is nil or empty")
1975+
}
1976+
1977+
sidebarTaskID := task.GetSidebarApp().GetId()
1978+
if !slices.Contains(appIDs, sidebarTaskID) {
1979+
warnUnknownSidebarAppID = true
19691980
}
19701981

1971-
id, err := uuid.Parse(task.SidebarApp.Id)
1982+
id, err := uuid.Parse(task.GetSidebarApp().GetId())
19721983
if err != nil {
19731984
return xerrors.Errorf("parse sidebar app id: %w", err)
19741985
}
19751986

19761987
sidebarAppID = uuid.NullUUID{UUID: id, Valid: true}
19771988
}
19781989

1990+
if warnUnknownSidebarAppID {
1991+
// Ref: https://github.com/coder/coder/issues/18776
1992+
// This can happen for a number of reasons:
1993+
// 1. Misconfigured template
1994+
// 2. Count=0 on the agent due to stop transition, meaning the associated coder_app was not inserted.
1995+
// Failing the build at this point is not ideal, so log a warning instead.
1996+
s.Logger.Warn(ctx, "unknown ai_task_sidebar_app_id",
1997+
slog.F("ai_task_sidebar_app_id", sidebarAppID.UUID.String()),
1998+
slog.F("job_id", job.ID.String()),
1999+
slog.F("workspace_id", workspace.ID),
2000+
slog.F("workspace_build_id", workspaceBuild.ID),
2001+
slog.F("transition", string(workspaceBuild.Transition)),
2002+
)
2003+
// In order to surface this to the user, we will also insert a warning into the build logs.
2004+
if _, err := db.InsertProvisionerJobLogs(ctx, database.InsertProvisionerJobLogsParams{
2005+
JobID: jobID,
2006+
CreatedAt: []time.Time{now, now, now, now},
2007+
Source: []database.LogSource{database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon, database.LogSourceProvisionerDaemon},
2008+
Level: []database.LogLevel{database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn, database.LogLevelWarn},
2009+
Stage: []string{"Cleaning Up", "Cleaning Up", "Cleaning Up", "Cleaning Up"},
2010+
Output: []string{
2011+
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()),
2012+
"Template author: double-check the following:",
2013+
" - 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).",
2014+
" - 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.",
2015+
},
2016+
}); err != nil {
2017+
s.Logger.Error(ctx, "insert provisioner job log for ai task sidebar app id warning",
2018+
slog.F("job_id", jobID),
2019+
slog.F("workspace_id", workspace.ID),
2020+
slog.F("workspace_build_id", workspaceBuild.ID),
2021+
slog.F("transition", string(workspaceBuild.Transition)),
2022+
)
2023+
}
2024+
// Important: reset hasAITask and sidebarAppID so that we don't run into a fk constraint violation.
2025+
hasAITask = false
2026+
sidebarAppID = uuid.NullUUID{}
2027+
}
2028+
19792029
// Regardless of whether there is an AI task or not, update the field to indicate one way or the other since it
19802030
// always defaults to nil. ONLY if has_ai_task=true MUST ai_task_sidebar_app_id be set.
1981-
err = db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
2031+
if err := db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
19822032
ID: workspaceBuild.ID,
19832033
HasAITask: sql.NullBool{
19842034
Bool: hasAITask,
19852035
Valid: true,
19862036
},
19872037
SidebarAppID: sidebarAppID,
19882038
UpdatedAt: now,
1989-
})
1990-
if err != nil {
2039+
}); err != nil {
19912040
return xerrors.Errorf("update workspace build ai tasks flag: %w", err)
19922041
}
19932042

‎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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy