Skip to content

Commit 9e64cdd

Browse files
mafredrimafredri
committed
improve tests and dbmem
1 parent c6e1cc6 commit 9e64cdd

File tree

3 files changed

+54
-31
lines changed

3 files changed

+54
-31
lines changed

‎cli/provisioners_test.go

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,20 +22,26 @@ func TestProvisioners(t *testing.T) {
2222

2323
db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
2424
client, _, coderdAPI := coderdtest.NewWithAPI(t, &coderdtest.Options{
25-
IncludeProvisionerDaemon: true,
25+
IncludeProvisionerDaemon: false,
2626
Database: db,
2727
Pubsub: ps,
2828
})
2929
owner := coderdtest.CreateFirstUser(t, client)
3030
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
3131

32+
// Create initial resources with a running provisioner.
33+
firstProvisioner := coderdtest.NewProvisionerDaemon(t, coderdAPI)
34+
defer firstProvisioner.Close()
3235
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, completeWithAgent())
3336
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
3437
template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
3538

3639
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
3740
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
3841

42+
// Stop the provisioner so it doesn't grab jobs.
43+
firstProvisioner.Close()
44+
3945
// Create a provisioner that's working on a job.
4046
pd1 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
4147
Name: "provisioner-1",
@@ -52,6 +58,7 @@ func TestProvisioners(t *testing.T) {
5258
WorkerID: uuid.NullUUID{UUID: pd1.ID, Valid: true},
5359
Input: json.RawMessage(`{"workspace_build_id":"` + wb1ID.String() + `"}`),
5460
StartedAt: sql.NullTime{Time: coderdAPI.Clock.Now(), Valid: true},
61+
Tags: database.StringMap{"owner": "", "scope": "organization"},
5562
})
5663
dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
5764
ID: wb1ID,
@@ -78,6 +85,7 @@ func TestProvisioners(t *testing.T) {
7885
Input: json.RawMessage(`{"workspace_build_id":"` + wb2ID.String() + `"}`),
7986
StartedAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-2 * time.Hour), Valid: true},
8087
CompletedAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Hour), Valid: true},
88+
Tags: database.StringMap{"owner": "", "scope": "organization"},
8189
})
8290
dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
8391
ID: wb2ID,
@@ -86,6 +94,24 @@ func TestProvisioners(t *testing.T) {
8694
TemplateVersionID: version.ID,
8795
})
8896

97+
// Create a pending job.
98+
w3 := dbgen.Workspace(t, coderdAPI.Database, database.WorkspaceTable{
99+
OwnerID: memberUser.ID,
100+
TemplateID: template.ID,
101+
})
102+
wb3ID := uuid.MustParse("00000000-0000-0000-bbbb-000000000003")
103+
job3 := dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
104+
ID: uuid.MustParse("00000000-0000-0000-cccc-000000000003"),
105+
Input: json.RawMessage(`{"workspace_build_id":"` + wb3ID.String() + `"}`),
106+
Tags: database.StringMap{"owner": "", "scope": "organization"},
107+
})
108+
dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
109+
ID: wb3ID,
110+
JobID: job3.ID,
111+
WorkspaceID: w3.ID,
112+
TemplateVersionID: version.ID,
113+
})
114+
89115
// Create a provisioner that is idle.
90116
pd3 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
91117
Name: "provisioner-3",

‎coderd/database/dbgen/dbgen.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -599,7 +599,6 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
599599
require.NoError(t, err)
600600
// There is no easy way to make sure we acquire the correct job.
601601
require.Equal(t, jobID, job.ID, "acquired incorrect job")
602-
fmt.Printf("%#v\n", job)
603602
}
604603

605604
if !orig.CompletedAt.Time.IsZero() || orig.Error.String != "" {

‎coderd/database/dbmem/dbmem.go

Lines changed: 27 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -3876,10 +3876,17 @@ func (q *FakeQuerier) GetProvisionerJobsByIDs(_ context.Context, ids []uuid.UUID
38763876
return jobs, nil
38773877
}
38783878

3879-
func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
3879+
func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(ctx context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
38803880
q.mutex.RLock()
38813881
defer q.mutex.RUnlock()
38823882

3883+
if ids == nil {
3884+
ids = []uuid.UUID{}
3885+
}
3886+
return q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, ids)
3887+
}
3888+
3889+
func (q *FakeQuerier) getProvisionerJobsByIDsWithQueuePositionLocked(_ context.Context, ids []uuid.UUID) ([]database.GetProvisionerJobsByIDsWithQueuePositionRow, error) {
38833890
// WITH pending_jobs AS (
38843891
// SELECT
38853892
// id, created_at
@@ -3948,7 +3955,7 @@ func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(_ context.Context
39483955
// pj.id IN (...)
39493956
jobs := make([]database.GetProvisionerJobsByIDsWithQueuePositionRow, 0)
39503957
for _, job := range q.provisionerJobs {
3951-
if !slices.Contains(ids, job.ID) {
3958+
if ids != nil && !slices.Contains(ids, job.ID) {
39523959
continue
39533960
}
39543961
// clone the Tags before appending, since maps are reference types and
@@ -3969,7 +3976,7 @@ func (q *FakeQuerier) GetProvisionerJobsByIDsWithQueuePosition(_ context.Context
39693976
return jobs, nil
39703977
}
39713978

3972-
func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(_ context.Context, arg database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) {
3979+
func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisioner(ctx context.Context, arg database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerParams) ([]database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow, error) {
39733980
err := validateDatabaseType(arg)
39743981
if err != nil {
39753982
return nil, err
@@ -4035,44 +4042,35 @@ func (q *FakeQuerier) GetProvisionerJobsByOrganizationAndStatusWithQueuePosition
40354042
LIMIT
40364043
sqlc.narg('limit')::int;
40374044
*/
4045+
rowsWithQueuePosition, err := q.getProvisionerJobsByIDsWithQueuePositionLocked(ctx, nil)
4046+
if err != nil {
4047+
return nil, err
4048+
}
4049+
40384050
var rows []database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow
4039-
var queuedJobs []database.ProvisionerJob
4040-
for _, job := range q.provisionerJobs {
4051+
for _, rowQP := range rowsWithQueuePosition {
4052+
job := rowQP.ProvisionerJob
4053+
40414054
if arg.OrganizationID.Valid && job.OrganizationID != arg.OrganizationID.UUID {
40424055
continue
40434056
}
40444057
if len(arg.Status) > 0 && !slices.Contains(arg.Status, job.JobStatus) {
40454058
continue
40464059
}
40474060

4048-
if job.JobStatus == database.ProvisionerJobStatusPending {
4049-
queuedJobs = append(queuedJobs, job)
4050-
}
40514061
row := database.GetProvisionerJobsByOrganizationAndStatusWithQueuePositionAndProvisionerRow{
4052-
ProvisionerJob: job,
4062+
ProvisionerJob: rowQP.ProvisionerJob,
4063+
QueuePosition: rowQP.QueuePosition,
4064+
QueueSize: rowQP.QueueSize,
40534065
}
4054-
rows = append(rows, row)
4055-
}
4056-
slices.SortFunc(queuedJobs, func(a, b database.ProvisionerJob) int {
4057-
return a.CreatedAt.Compare(b.CreatedAt)
4058-
})
4059-
for i := range rows {
4060-
queuePosition := slices.IndexFunc(queuedJobs, func(job database.ProvisionerJob) bool {
4061-
return job.ID == rows[i].ProvisionerJob.ID
4062-
})
4063-
if queuePosition >= 0 {
4064-
rows[i].QueuePosition = int64(queuePosition + 1)
4065-
rows[i].QueueSize = int64(len(queuedJobs))
4066-
4067-
for _, daemon := range q.provisionerDaemons {
4068-
job := rows[i].ProvisionerJob
4069-
if daemon.OrganizationID == job.OrganizationID &&
4070-
slices.Contains(daemon.Provisioners, job.Provisioner) &&
4071-
tagsSubset(job.Tags, daemon.Tags) {
4072-
rows[i].AvailableWorkers = append(rows[i].AvailableWorkers, daemon.ID)
4073-
}
4066+
for _, daemon := range q.provisionerDaemons {
4067+
if daemon.OrganizationID == job.OrganizationID &&
4068+
slices.Contains(daemon.Provisioners, job.Provisioner) &&
4069+
tagsSubset(job.Tags, daemon.Tags) {
4070+
row.AvailableWorkers = append(row.AvailableWorkers, daemon.ID)
40744071
}
40754072
}
4073+
rows = append(rows, row)
40764074
}
40774075

40784076
return rows, nil

0 commit comments

Comments
 (0)