Skip to content

Commit d3ae1c3

Browse files
mafredrimafredri
committed
test
1 parent d6248d7 commit d3ae1c3

File tree

2 files changed

+178
-32
lines changed

2 files changed

+178
-32
lines changed

‎cli/provisioners_test.go

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
package cli_test
2+
3+
import (
4+
"database/sql"
5+
"encoding/json"
6+
"testing"
7+
"time"
8+
9+
"github.com/google/uuid"
10+
"github.com/stretchr/testify/require"
11+
12+
"github.com/coder/coder/v2/cli/clitest"
13+
"github.com/coder/coder/v2/coderd/coderdtest"
14+
"github.com/coder/coder/v2/coderd/database"
15+
"github.com/coder/coder/v2/coderd/database/dbgen"
16+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
17+
"github.com/coder/coder/v2/codersdk"
18+
)
19+
20+
func TestProvisioners(t *testing.T) {
21+
t.Parallel()
22+
23+
db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
24+
client, _, coderdAPI := coderdtest.NewWithAPI(t, &coderdtest.Options{
25+
IncludeProvisionerDaemon: true,
26+
Database: db,
27+
Pubsub: ps,
28+
})
29+
owner := coderdtest.CreateFirstUser(t, client)
30+
member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
31+
32+
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, completeWithAgent())
33+
coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
34+
template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
35+
36+
workspace := coderdtest.CreateWorkspace(t, client, template.ID)
37+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
38+
39+
// Create a provisioner that's working on a job.
40+
pd1 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
41+
Name: "provisioner-1",
42+
CreatedAt: timeParse(t, "2006-01-02", "2024-12-20"),
43+
KeyID: uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
44+
})
45+
w1 := dbgen.Workspace(t, coderdAPI.Database, database.WorkspaceTable{
46+
OwnerID: memberUser.ID,
47+
TemplateID: template.ID,
48+
})
49+
wb1ID := uuid.MustParse("00000000-0000-0000-bbbb-000000000001")
50+
job1 := dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
51+
ID: uuid.MustParse("00000000-0000-0000-cccc-000000000001"),
52+
WorkerID: uuid.NullUUID{UUID: pd1.ID, Valid: true},
53+
Input: json.RawMessage(`{"workspace_build_id":"` + wb1ID.String() + `"}`),
54+
StartedAt: sql.NullTime{Time: coderdAPI.Clock.Now(), Valid: true},
55+
})
56+
dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
57+
ID: wb1ID,
58+
JobID: job1.ID,
59+
WorkspaceID: w1.ID,
60+
TemplateVersionID: version.ID,
61+
})
62+
63+
// Create another provisioner that completed a job and is offline.
64+
pd2 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
65+
Name: "provisioner-2",
66+
CreatedAt: timeParse(t, "2006-01-02", "2024-12-20"),
67+
LastSeenAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Hour), Valid: true},
68+
KeyID: uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
69+
})
70+
w2 := dbgen.Workspace(t, coderdAPI.Database, database.WorkspaceTable{
71+
OwnerID: memberUser.ID,
72+
TemplateID: template.ID,
73+
})
74+
wb2ID := uuid.MustParse("00000000-0000-0000-bbbb-000000000002")
75+
job2 := dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
76+
ID: uuid.MustParse("00000000-0000-0000-cccc-000000000002"),
77+
WorkerID: uuid.NullUUID{UUID: pd2.ID, Valid: true},
78+
Input: json.RawMessage(`{"workspace_build_id":"` + wb2ID.String() + `"}`),
79+
StartedAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-2 * time.Hour), Valid: true},
80+
CompletedAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Hour), Valid: true},
81+
})
82+
dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
83+
ID: wb2ID,
84+
JobID: job2.ID,
85+
WorkspaceID: w2.ID,
86+
TemplateVersionID: version.ID,
87+
})
88+
89+
// Create a provisioner that is idle.
90+
pd3 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
91+
Name: "provisioner-3",
92+
CreatedAt: timeParse(t, "2006-01-02", "2024-12-20"),
93+
KeyID: uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
94+
})
95+
_ = pd3
96+
97+
t.Run("list", func(t *testing.T) {
98+
t.Parallel()
99+
100+
inv, root := clitest.New(t,
101+
"provisioners",
102+
"list",
103+
"--column", "id,created at,last seen at,name,version,api version,tags,status,current job id,previous job id,previous job status,organization",
104+
)
105+
clitest.SetupConfig(t, member, root)
106+
err := inv.Run()
107+
require.NoError(t, err)
108+
109+
// TODO(mafredri): Verify golden output.
110+
})
111+
112+
t.Run("jobs list", func(t *testing.T) {
113+
t.Parallel()
114+
115+
inv, root := clitest.New(t,
116+
"provisioners",
117+
"jobs",
118+
"list",
119+
"--column", "id,created at,status,worker id,tags,template version id,workspace build id,type,available workers,organization,queue",
120+
)
121+
clitest.SetupConfig(t, member, root)
122+
err := inv.Run()
123+
require.NoError(t, err)
124+
125+
// TODO(mafredri): Verify golden output.
126+
})
127+
}
128+
129+
func timeParse(t *testing.T, layout, s string) time.Time {
130+
t.Helper()
131+
tm, err := time.Parse(layout, s)
132+
require.NoError(t, err)
133+
return tm
134+
}

‎coderd/database/dbgen/dbgen.go

Lines changed: 44 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -248,12 +248,18 @@ func WorkspaceAgentScriptTiming(t testing.TB, db database.Store, orig database.W
248248
func Workspace(t testing.TB, db database.Store, orig database.WorkspaceTable) database.WorkspaceTable {
249249
t.Helper()
250250

251+
var defOrgID uuid.UUID
252+
if orig.OrganizationID == uuid.Nil {
253+
defOrg, _ := db.GetDefaultOrganization(genCtx)
254+
defOrgID = defOrg.ID
255+
}
256+
251257
workspace, err := db.InsertWorkspace(genCtx, database.InsertWorkspaceParams{
252258
ID: takeFirst(orig.ID, uuid.New()),
253259
OwnerID: takeFirst(orig.OwnerID, uuid.New()),
254260
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
255261
UpdatedAt: takeFirst(orig.UpdatedAt, dbtime.Now()),
256-
OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
262+
OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
257263
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
258264
LastUsedAt: takeFirst(orig.LastUsedAt, dbtime.Now()),
259265
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
@@ -505,9 +511,27 @@ func GroupMember(t testing.TB, db database.Store, member database.GroupMemberTab
505511

506512
// ProvisionerDaemon creates a provisioner daemon as far as the database is concerned. It does not run a provisioner daemon.
507513
// If no key is provided, it will create one.
508-
func ProvisionerDaemon(t testing.TB, db database.Store, daemon database.ProvisionerDaemon) database.ProvisionerDaemon {
514+
func ProvisionerDaemon(t testing.TB, db database.Store, orig database.ProvisionerDaemon) database.ProvisionerDaemon {
509515
t.Helper()
510516

517+
var defOrgID uuid.UUID
518+
if orig.OrganizationID == uuid.Nil {
519+
defOrg, _ := db.GetDefaultOrganization(genCtx)
520+
defOrgID = defOrg.ID
521+
}
522+
523+
daemon := database.UpsertProvisionerDaemonParams{
524+
Name: takeFirst(orig.Name, testutil.GetRandomName(t)),
525+
OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
526+
CreatedAt: takeFirst(orig.CreatedAt, dbtime.Now()),
527+
Provisioners: takeFirstSlice(orig.Provisioners, []database.ProvisionerType{database.ProvisionerTypeEcho}),
528+
Tags: takeFirstMap(orig.Tags, database.StringMap{"owner": "", "scope": "organization"}),
529+
KeyID: takeFirst(orig.KeyID, uuid.Nil),
530+
LastSeenAt: takeFirst(orig.LastSeenAt, sql.NullTime{Time: dbtime.Now(), Valid: true}),
531+
Version: takeFirst(orig.Version, "v0.0.0"),
532+
APIVersion: takeFirst(orig.APIVersion, "1.1"),
533+
}
534+
511535
if daemon.KeyID == uuid.Nil {
512536
key, err := db.InsertProvisionerKey(genCtx, database.InsertProvisionerKeyParams{
513537
ID: uuid.New(),
@@ -521,24 +545,7 @@ func ProvisionerDaemon(t testing.TB, db database.Store, daemon database.Provisio
521545
daemon.KeyID = key.ID
522546
}
523547

524-
if daemon.CreatedAt.IsZero() {
525-
daemon.CreatedAt = dbtime.Now()
526-
}
527-
if daemon.Name == "" {
528-
daemon.Name = "test-daemon"
529-
}
530-
531-
d, err := db.UpsertProvisionerDaemon(genCtx, database.UpsertProvisionerDaemonParams{
532-
Name: daemon.Name,
533-
OrganizationID: daemon.OrganizationID,
534-
CreatedAt: daemon.CreatedAt,
535-
Provisioners: daemon.Provisioners,
536-
Tags: daemon.Tags,
537-
KeyID: daemon.KeyID,
538-
LastSeenAt: daemon.LastSeenAt,
539-
Version: daemon.Version,
540-
APIVersion: daemon.APIVersion,
541-
})
548+
d, err := db.UpsertProvisionerDaemon(genCtx, daemon)
542549
require.NoError(t, err)
543550
return d
544551
}
@@ -556,12 +563,10 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
556563

557564
jobID := takeFirst(orig.ID, uuid.New())
558565
// Always set some tags to prevent Acquire from grabbing jobs it should not.
559-
if !orig.StartedAt.Time.IsZero() {
560-
if orig.Tags == nil {
561-
orig.Tags = make(database.StringMap)
562-
}
566+
tags := takeFirstMap(orig.Tags, database.StringMap{"owner": "", "scope": "organization"})
567+
if orig.Tags == nil && !orig.StartedAt.Time.IsZero() {
563568
// Make sure when we acquire the job, we only get this one.
564-
orig.Tags[jobID.String()] = "true"
569+
tags[jobID.String()] = "true"
565570
}
566571

567572
job, err := db.InsertProvisionerJob(genCtx, database.InsertProvisionerJobParams{
@@ -575,7 +580,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
575580
FileID: takeFirst(orig.FileID, uuid.New()),
576581
Type: takeFirst(orig.Type, database.ProvisionerJobTypeWorkspaceBuild),
577582
Input: takeFirstSlice(orig.Input, []byte("{}")),
578-
Tags: orig.Tags,
583+
Tags: tags,
579584
TraceMetadata: pqtype.NullRawMessage{},
580585
})
581586
require.NoError(t, err, "insert job")
@@ -587,17 +592,18 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
587592
job, err = db.AcquireProvisionerJob(genCtx, database.AcquireProvisionerJobParams{
588593
StartedAt: orig.StartedAt,
589594
OrganizationID: job.OrganizationID,
590-
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
591-
ProvisionerTags: must(json.Marshal(orig.Tags)),
592-
WorkerID: uuid.NullUUID{},
595+
Types: []database.ProvisionerType{job.Provisioner},
596+
ProvisionerTags: must(json.Marshal(tags)),
597+
WorkerID: takeFirst(orig.WorkerID, uuid.NullUUID{}),
593598
})
594599
require.NoError(t, err)
595600
// There is no easy way to make sure we acquire the correct job.
596601
require.Equal(t, jobID, job.ID, "acquired incorrect job")
602+
fmt.Printf("%#v\n", job)
597603
}
598604

599605
if !orig.CompletedAt.Time.IsZero() || orig.Error.String != "" {
600-
err := db.UpdateProvisionerJobWithCompleteByID(genCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
606+
err = db.UpdateProvisionerJobWithCompleteByID(genCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
601607
ID: jobID,
602608
UpdatedAt: job.UpdatedAt,
603609
CompletedAt: orig.CompletedAt,
@@ -607,7 +613,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
607613
require.NoError(t, err)
608614
}
609615
if !orig.CanceledAt.Time.IsZero() {
610-
err := db.UpdateProvisionerJobWithCancelByID(genCtx, database.UpdateProvisionerJobWithCancelByIDParams{
616+
err = db.UpdateProvisionerJobWithCancelByID(genCtx, database.UpdateProvisionerJobWithCancelByIDParams{
611617
ID: jobID,
612618
CanceledAt: orig.CanceledAt,
613619
CompletedAt: orig.CompletedAt,
@@ -616,7 +622,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
616622
}
617623

618624
job, err = db.GetProvisionerJobByID(genCtx, jobID)
619-
require.NoError(t, err)
625+
require.NoError(t, err, "get job: %s", jobID.String())
620626

621627
return job
622628
}
@@ -1108,6 +1114,12 @@ func takeFirstSlice[T any](values ...[]T) []T {
11081114
})
11091115
}
11101116

1117+
func takeFirstMap[T, E comparable](values ...map[T]E) map[T]E {
1118+
return takeFirstF(values, func(v map[T]E) bool {
1119+
return v != nil
1120+
})
1121+
}
1122+
11111123
// takeFirstF takes the first value that returns true
11121124
func takeFirstF[Value any](values []Value, take func(v Value) bool) Value {
11131125
for _, v := range values {

0 commit comments

Comments
 (0)