Skip to content

Commit b1c5a08

Browse files
kylecarbskylecarbs
committed
feat: add dbfakedata for workspace builds and resources
This creates `coderdtest.NewWithDatabase` and adds a series of helper functions to `dbfake` that insert structured fake data for resources into the database. It allows us to remove provisionerd from a significant amount of tests which should speed them up and reduce flakes.
1 parent 6b7858c commit b1c5a08

File tree

5 files changed

+143
-81
lines changed

5 files changed

+143
-81
lines changed

‎cli/agent_test.go

+30-48
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ import (
1616
"github.com/coder/coder/v2/agent"
1717
"github.com/coder/coder/v2/cli/clitest"
1818
"github.com/coder/coder/v2/coderd/coderdtest"
19+
"github.com/coder/coder/v2/coderd/database"
20+
"github.com/coder/coder/v2/coderd/database/dbfakedata"
1921
"github.com/coder/coder/v2/codersdk"
2022
"github.com/coder/coder/v2/provisioner/echo"
2123
"github.com/coder/coder/v2/provisionersdk/proto"
@@ -68,33 +70,23 @@ func TestWorkspaceAgent(t *testing.T) {
6870
t.Parallel()
6971
instanceID := "instanceidentifier"
7072
certificates, metadataClient := coderdtest.NewAzureInstanceIdentity(t, instanceID)
71-
client := coderdtest.New(t, &coderdtest.Options{
72-
AzureCertificates: certificates,
73-
IncludeProvisionerDaemon: true,
73+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
74+
AzureCertificates: certificates,
7475
})
7576
user := coderdtest.CreateFirstUser(t, client)
76-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
77-
Parse: echo.ParseComplete,
78-
ProvisionApply: []*proto.Response{{
79-
Type: &proto.Response_Apply{
80-
Apply: &proto.ApplyComplete{
81-
Resources: []*proto.Resource{{
82-
Name: "somename",
83-
Type: "someinstance",
84-
Agents: []*proto.Agent{{
85-
Auth: &proto.Agent_InstanceId{
86-
InstanceId: instanceID,
87-
},
88-
}},
89-
}},
90-
},
77+
ws := dbfakedata.CreateWorkspace(t, db, database.Workspace{
78+
OrganizationID: user.OrganizationID,
79+
OwnerID: user.UserID,
80+
})
81+
dbfakedata.CreateWorkspaceBuild(t, db, ws, database.WorkspaceBuild{}, &proto.Resource{
82+
Name: "somename",
83+
Type: "someinstance",
84+
Agents: []*proto.Agent{{
85+
Auth: &proto.Agent_InstanceId{
86+
InstanceId: instanceID,
9187
},
9288
}},
9389
})
94-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
95-
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
96-
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
97-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
9890

9991
inv, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String())
10092
inv = inv.WithContext(
@@ -103,8 +95,8 @@ func TestWorkspaceAgent(t *testing.T) {
10395
)
10496
ctx := inv.Context()
10597
clitest.Start(t, inv)
106-
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
107-
workspace, err := client.Workspace(ctx, workspace.ID)
98+
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
99+
workspace, err := client.Workspace(ctx, ws.ID)
108100
require.NoError(t, err)
109101
resources := workspace.LatestBuild.Resources
110102
if assert.NotEmpty(t, workspace.LatestBuild.Resources) && assert.NotEmpty(t, resources[0].Agents) {
@@ -120,33 +112,23 @@ func TestWorkspaceAgent(t *testing.T) {
120112
t.Parallel()
121113
instanceID := "instanceidentifier"
122114
certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
123-
client := coderdtest.New(t, &coderdtest.Options{
124-
AWSCertificates: certificates,
125-
IncludeProvisionerDaemon: true,
115+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
116+
AWSCertificates: certificates,
126117
})
127118
user := coderdtest.CreateFirstUser(t, client)
128-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
129-
Parse: echo.ParseComplete,
130-
ProvisionApply: []*proto.Response{{
131-
Type: &proto.Response_Apply{
132-
Apply: &proto.ApplyComplete{
133-
Resources: []*proto.Resource{{
134-
Name: "somename",
135-
Type: "someinstance",
136-
Agents: []*proto.Agent{{
137-
Auth: &proto.Agent_InstanceId{
138-
InstanceId: instanceID,
139-
},
140-
}},
141-
}},
142-
},
119+
ws := dbfakedata.CreateWorkspace(t, db, database.Workspace{
120+
OrganizationID: user.OrganizationID,
121+
OwnerID: user.UserID,
122+
})
123+
dbfakedata.CreateWorkspaceBuild(t, db, ws, database.WorkspaceBuild{}, &proto.Resource{
124+
Name: "somename",
125+
Type: "someinstance",
126+
Agents: []*proto.Agent{{
127+
Auth: &proto.Agent_InstanceId{
128+
InstanceId: instanceID,
143129
},
144130
}},
145131
})
146-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
147-
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
148-
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
149-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
150132

151133
inv, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String())
152134
inv = inv.WithContext(
@@ -155,8 +137,8 @@ func TestWorkspaceAgent(t *testing.T) {
155137
)
156138
clitest.Start(t, inv)
157139
ctx := inv.Context()
158-
coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
159-
workspace, err := client.Workspace(ctx, workspace.ID)
140+
coderdtest.AwaitWorkspaceAgents(t, client, ws.ID)
141+
workspace, err := client.Workspace(ctx, ws.ID)
160142
require.NoError(t, err)
161143
resources := workspace.LatestBuild.Resources
162144
if assert.NotEmpty(t, resources) && assert.NotEmpty(t, resources[0].Agents) {

‎coderd/coderdtest/coderdtest.go

+7
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,13 @@ func New(t testing.TB, options *Options) *codersdk.Client {
151151
return client
152152
}
153153

154+
// NewWithDatabase constructs a codersdk client connected to an in-memory API instance.
155+
// The database is returned to provide direct data manipulation for tests.
156+
func NewWithDatabase(t testing.TB, options *Options) (*codersdk.Client, database.Store) {
157+
client, _, api := NewWithAPI(t, options)
158+
return client, api.Database
159+
}
160+
154161
// NewWithProvisionerCloser returns a client as well as a handle to close
155162
// the provisioner. This is a temporary function while work is done to
156163
// standardize how provisioners are registered with coderd. The option

‎coderd/database/dbfakedata/dbfakedata.go

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
package dbfakedata
2+
3+
import (
4+
"context"
5+
"database/sql"
6+
"encoding/json"
7+
"testing"
8+
9+
"github.com/google/uuid"
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/coder/coder/v2/coderd/database"
13+
"github.com/coder/coder/v2/coderd/database/dbauthz"
14+
"github.com/coder/coder/v2/coderd/database/dbgen"
15+
"github.com/coder/coder/v2/coderd/database/dbtime"
16+
"github.com/coder/coder/v2/coderd/provisionerdserver"
17+
"github.com/coder/coder/v2/coderd/telemetry"
18+
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
19+
)
20+
21+
// CreateWorkspace inserts a workspace into the database.
22+
func CreateWorkspace(t testing.TB, db database.Store, seed database.Workspace) database.Workspace {
23+
t.Helper()
24+
25+
// This intentionally fulfills the minimum requirements of the schema.
26+
// Tests can provide a custom template ID if necessary.
27+
if seed.TemplateID == uuid.Nil {
28+
template := dbgen.Template(t, db, database.Template{
29+
OrganizationID: seed.OrganizationID,
30+
CreatedBy: seed.OwnerID,
31+
})
32+
seed.TemplateID = template.ID
33+
}
34+
return dbgen.Workspace(t, db, seed)
35+
}
36+
37+
// CreateWorkspaceBuild inserts a build and a successful job into the database.
38+
func CreateWorkspaceBuild(t testing.TB, db database.Store, ws database.Workspace, seed database.WorkspaceBuild, resources ...*sdkproto.Resource) database.WorkspaceBuild {
39+
t.Helper()
40+
jobID := uuid.New()
41+
seed.JobID = jobID
42+
seed.WorkspaceID = ws.ID
43+
// This intentionally fulfills the minimum requirements of the schema.
44+
// Tests can provide a custom version ID if necessary.
45+
if seed.TemplateVersionID == uuid.Nil {
46+
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
47+
OrganizationID: ws.OrganizationID,
48+
CreatedBy: ws.OwnerID,
49+
})
50+
seed.TemplateVersionID = templateVersion.ID
51+
}
52+
build := dbgen.WorkspaceBuild(t, db, seed)
53+
54+
payload, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
55+
WorkspaceBuildID: build.ID,
56+
})
57+
require.NoError(t, err)
58+
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
59+
ID: jobID,
60+
Input: payload,
61+
OrganizationID: ws.OrganizationID,
62+
CompletedAt: sql.NullTime{
63+
Time: dbtime.Now(),
64+
Valid: true,
65+
},
66+
})
67+
CreateProvisionerJobResources(t, db, job.ID, seed.Transition, resources...)
68+
return build
69+
}
70+
71+
// CreateProvisionerJobResources inserts a series of resources into a provisioner job.
72+
func CreateProvisionerJobResources(t testing.TB, db database.Store, job uuid.UUID, transition database.WorkspaceTransition, resources ...*sdkproto.Resource) {
73+
t.Helper()
74+
if transition == "" {
75+
// Default to start!
76+
transition = database.WorkspaceTransitionStart
77+
}
78+
for _, resource := range resources {
79+
//nolint:gocritic // This is only used by tests.
80+
err := provisionerdserver.InsertWorkspaceResource(dbauthz.AsSystemRestricted(context.Background()), db, job, transition, resource, &telemetry.Snapshot{})
81+
require.NoError(t, err)
82+
}
83+
}

‎coderd/database/dbgen/dbgen.go

+7-1
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ func AuditLog(t testing.TB, db database.Store, seed database.AuditLog) database.
6666

6767
func Template(t testing.TB, db database.Store, seed database.Template) database.Template {
6868
id := takeFirst(seed.ID, uuid.New())
69+
if seed.GroupACL == nil {
70+
// By default, all users in the organization can read the template.
71+
seed.GroupACL = database.TemplateACL{
72+
seed.OrganizationID.String(): []rbac.Action{rbac.ActionRead},
73+
}
74+
}
6975
err := db.InsertTemplate(genCtx, database.InsertTemplateParams{
7076
ID: id,
7177
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
@@ -84,7 +90,7 @@ func Template(t testing.TB, db database.Store, seed database.Template) database.
8490
})
8591
require.NoError(t, err, "insert template")
8692

87-
template, err := db.GetTemplateByID(context.Background(), id)
93+
template, err := db.GetTemplateByID(genCtx, id)
8894
require.NoError(t, err, "get template")
8995
return template
9096
}

‎coderd/workspaceagents_test.go

+16-32
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"github.com/coder/coder/v2/agent/agenttest"
2424
"github.com/coder/coder/v2/coderd/coderdtest"
2525
"github.com/coder/coder/v2/coderd/database"
26+
"github.com/coder/coder/v2/coderd/database/dbfakedata"
2627
"github.com/coder/coder/v2/coderd/database/dbtime"
2728
"github.com/coder/coder/v2/codersdk"
2829
"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -36,44 +37,27 @@ func TestWorkspaceAgent(t *testing.T) {
3637
t.Parallel()
3738
t.Run("Connect", func(t *testing.T) {
3839
t.Parallel()
39-
client := coderdtest.New(t, &coderdtest.Options{
40-
IncludeProvisionerDaemon: true,
41-
})
40+
client, db := coderdtest.NewWithDatabase(t, nil)
4241
user := coderdtest.CreateFirstUser(t, client)
43-
authToken := uuid.NewString()
4442
tmpDir := t.TempDir()
45-
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
46-
Parse: echo.ParseComplete,
47-
ProvisionPlan: echo.PlanComplete,
48-
ProvisionApply: []*proto.Response{{
49-
Type: &proto.Response_Apply{
50-
Apply: &proto.ApplyComplete{
51-
Resources: []*proto.Resource{{
52-
Name: "example",
53-
Type: "aws_instance",
54-
Agents: []*proto.Agent{{
55-
Id: uuid.NewString(),
56-
Directory: tmpDir,
57-
Auth: &proto.Agent_Token{
58-
Token: authToken,
59-
},
60-
}},
61-
}},
62-
},
63-
},
64-
}},
65-
})
66-
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
67-
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
68-
69-
anotherClient, _ := coderdtest.CreateAnotherUser(t, client, user.OrganizationID)
70-
workspace := coderdtest.CreateWorkspace(t, anotherClient, user.OrganizationID, template.ID)
71-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, anotherClient, workspace.LatestBuild.ID)
43+
anotherClient, anotherUser := coderdtest.CreateAnotherUser(t, client, user.OrganizationID)
7244

7345
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
7446
defer cancel()
7547

76-
workspace, err := anotherClient.Workspace(ctx, workspace.ID)
48+
ws := dbfakedata.CreateWorkspace(t, db, database.Workspace{
49+
OrganizationID: user.OrganizationID,
50+
OwnerID: anotherUser.ID,
51+
})
52+
dbfakedata.CreateWorkspaceBuild(t, db, ws, database.WorkspaceBuild{}, &proto.Resource{
53+
Name: "aws_instance",
54+
Agents: []*proto.Agent{{
55+
Id: uuid.NewString(),
56+
Name: "hey",
57+
Directory: tmpDir,
58+
}},
59+
})
60+
workspace, err := anotherClient.Workspace(ctx, ws.ID)
7761
require.NoError(t, err)
7862
require.Equal(t, tmpDir, workspace.LatestBuild.Resources[0].Agents[0].Directory)
7963
_, err = anotherClient.WorkspaceAgent(ctx, workspace.LatestBuild.Resources[0].Agents[0].ID)

0 commit comments

Comments
 (0)