test

coder · mafredri · Dec 19, 2024 · Dec 19, 2024 · Dec 19, 2024 · Dec 19, 2024
commit 8975d4160615ee04c9314ceb6a570bae7844dcd5
diff --git a/cli/provisioners_test.go b/cli/provisioners_test.go
@@ -0,0 +1,134 @@
+package cli_test
+
+import (
+	"database/sql"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+func TestProvisioners(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
+	client, _, coderdAPI := coderdtest.NewWithAPI(t, &coderdtest.Options{
+		IncludeProvisionerDaemon: true,
+		Database:                 db,
+		Pubsub:                   ps,
+	})
+	owner := coderdtest.CreateFirstUser(t, client)
+	member, memberUser := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
+
+	version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, completeWithAgent())
+	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+	template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
+
+	workspace := coderdtest.CreateWorkspace(t, client, template.ID)
+	coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+
+	// Create a provisioner that's working on a job.
+	pd1 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
+		Name:      "provisioner-1",
+		CreatedAt: timeParse(t, "2006-01-02", "2024-12-20"),
+		KeyID:     uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
+	})
+	w1 := dbgen.Workspace(t, coderdAPI.Database, database.WorkspaceTable{
+		OwnerID:    memberUser.ID,
+		TemplateID: template.ID,
+	})
+	wb1ID := uuid.MustParse("00000000-0000-0000-bbbb-000000000001")
+	job1 := dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
+		ID:        uuid.MustParse("00000000-0000-0000-cccc-000000000001"),
+		WorkerID:  uuid.NullUUID{UUID: pd1.ID, Valid: true},
+		Input:     json.RawMessage(`{"workspace_build_id":"` + wb1ID.String() + `"}`),
+		StartedAt: sql.NullTime{Time: coderdAPI.Clock.Now(), Valid: true},
+	})
+	dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
+		ID:                wb1ID,
+		JobID:             job1.ID,
+		WorkspaceID:       w1.ID,
+		TemplateVersionID: version.ID,
+	})
+
+	// Create another provisioner that completed a job and is offline.
+	pd2 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
+		Name:       "provisioner-2",
+		CreatedAt:  timeParse(t, "2006-01-02", "2024-12-20"),
+		LastSeenAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Hour), Valid: true},
+		KeyID:      uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
+	})
+	w2 := dbgen.Workspace(t, coderdAPI.Database, database.WorkspaceTable{
+		OwnerID:    memberUser.ID,
+		TemplateID: template.ID,
+	})
+	wb2ID := uuid.MustParse("00000000-0000-0000-bbbb-000000000002")
+	job2 := dbgen.ProvisionerJob(t, db, coderdAPI.Pubsub, database.ProvisionerJob{
+		ID:          uuid.MustParse("00000000-0000-0000-cccc-000000000002"),
+		WorkerID:    uuid.NullUUID{UUID: pd2.ID, Valid: true},
+		Input:       json.RawMessage(`{"workspace_build_id":"` + wb2ID.String() + `"}`),
+		StartedAt:   sql.NullTime{Time: coderdAPI.Clock.Now().Add(-2 * time.Hour), Valid: true},
+		CompletedAt: sql.NullTime{Time: coderdAPI.Clock.Now().Add(-time.Hour), Valid: true},
+	})
+	dbgen.WorkspaceBuild(t, coderdAPI.Database, database.WorkspaceBuild{
+		ID:                wb2ID,
+		JobID:             job2.ID,
+		WorkspaceID:       w2.ID,
+		TemplateVersionID: version.ID,
+	})
+
+	// Create a provisioner that is idle.
+	pd3 := dbgen.ProvisionerDaemon(t, coderdAPI.Database, database.ProvisionerDaemon{
+		Name:      "provisioner-3",
+		CreatedAt: timeParse(t, "2006-01-02", "2024-12-20"),
+		KeyID:     uuid.MustParse(codersdk.ProvisionerKeyIDBuiltIn),
+	})
+	_ = pd3
+
+	t.Run("list", func(t *testing.T) {
+		t.Parallel()
+
+		inv, root := clitest.New(t,
+			"provisioners",
+			"list",
+			"--column", "id,created at,last seen at,name,version,api version,tags,status,current job id,previous job id,previous job status,organization",
+		)
+		clitest.SetupConfig(t, member, root)
+		err := inv.Run()
+		require.NoError(t, err)
+
+		// TODO(mafredri): Verify golden output.
+	})
+
+	t.Run("jobs list", func(t *testing.T) {
+		t.Parallel()
+
+		inv, root := clitest.New(t,
+			"provisioners",
+			"jobs",
+			"list",
+			"--column", "id,created at,status,worker id,tags,template version id,workspace build id,type,available workers,organization,queue",
+		)
+		clitest.SetupConfig(t, member, root)
+		err := inv.Run()
+		require.NoError(t, err)
+
+		// TODO(mafredri): Verify golden output.
+	})
+}
+
+func timeParse(t *testing.T, layout, s string) time.Time {
+	t.Helper()
+	tm, err := time.Parse(layout, s)
+	require.NoError(t, err)
+	return tm
+}
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -248,12 +248,18 @@ func WorkspaceAgentScriptTiming(t testing.TB, db database.Store, orig database.W
 func Workspace(t testing.TB, db database.Store, orig database.WorkspaceTable) database.WorkspaceTable {
 	t.Helper()
 
+	var defOrgID uuid.UUID
+	if orig.OrganizationID == uuid.Nil {
+		defOrg, _ := db.GetDefaultOrganization(genCtx)
+		defOrgID = defOrg.ID
+	}
+
 	workspace, err := db.InsertWorkspace(genCtx, database.InsertWorkspaceParams{
 		ID:                takeFirst(orig.ID, uuid.New()),
 		OwnerID:           takeFirst(orig.OwnerID, uuid.New()),
 		CreatedAt:         takeFirst(orig.CreatedAt, dbtime.Now()),
 		UpdatedAt:         takeFirst(orig.UpdatedAt, dbtime.Now()),
-		OrganizationID:    takeFirst(orig.OrganizationID, uuid.New()),
+		OrganizationID:    takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
 		TemplateID:        takeFirst(orig.TemplateID, uuid.New()),
 		LastUsedAt:        takeFirst(orig.LastUsedAt, dbtime.Now()),
 		Name:              takeFirst(orig.Name, testutil.GetRandomName(t)),
@@ -505,9 +511,27 @@ func GroupMember(t testing.TB, db database.Store, member database.GroupMemberTab
 
 // ProvisionerDaemon creates a provisioner daemon as far as the database is concerned. It does not run a provisioner daemon.
 // If no key is provided, it will create one.
-func ProvisionerDaemon(t testing.TB, db database.Store, daemon database.ProvisionerDaemon) database.ProvisionerDaemon {
+func ProvisionerDaemon(t testing.TB, db database.Store, orig database.ProvisionerDaemon) database.ProvisionerDaemon {
 	t.Helper()
 
+	var defOrgID uuid.UUID
+	if orig.OrganizationID == uuid.Nil {
+		defOrg, _ := db.GetDefaultOrganization(genCtx)
+		defOrgID = defOrg.ID
+	}
+
+	daemon := database.UpsertProvisionerDaemonParams{
+		Name:           takeFirst(orig.Name, testutil.GetRandomName(t)),
+		OrganizationID: takeFirst(orig.OrganizationID, defOrgID, uuid.New()),
+		CreatedAt:      takeFirst(orig.CreatedAt, dbtime.Now()),
+		Provisioners:   takeFirstSlice(orig.Provisioners, []database.ProvisionerType{database.ProvisionerTypeEcho}),
+		Tags:           takeFirstMap(orig.Tags, database.StringMap{"owner": "", "scope": "organization"}),
+		KeyID:          takeFirst(orig.KeyID, uuid.Nil),
+		LastSeenAt:     takeFirst(orig.LastSeenAt, sql.NullTime{Time: dbtime.Now(), Valid: true}),
+		Version:        takeFirst(orig.Version, "v0.0.0"),
+		APIVersion:     takeFirst(orig.APIVersion, "1.1"),
+	}
+
 	if daemon.KeyID == uuid.Nil {
 		key, err := db.InsertProvisionerKey(genCtx, database.InsertProvisionerKeyParams{
 			ID:             uuid.New(),
@@ -521,24 +545,7 @@ func ProvisionerDaemon(t testing.TB, db database.Store, daemon database.Provisio
 		daemon.KeyID = key.ID
 	}
 
-	if daemon.CreatedAt.IsZero() {
-		daemon.CreatedAt = dbtime.Now()
-	}
-	if daemon.Name == "" {
-		daemon.Name = "test-daemon"
-	}
-
-	d, err := db.UpsertProvisionerDaemon(genCtx, database.UpsertProvisionerDaemonParams{
-		Name:           daemon.Name,
-		OrganizationID: daemon.OrganizationID,
-		CreatedAt:      daemon.CreatedAt,
-		Provisioners:   daemon.Provisioners,
-		Tags:           daemon.Tags,
-		KeyID:          daemon.KeyID,
-		LastSeenAt:     daemon.LastSeenAt,
-		Version:        daemon.Version,
-		APIVersion:     daemon.APIVersion,
-	})
+	d, err := db.UpsertProvisionerDaemon(genCtx, daemon)
 	require.NoError(t, err)
 	return d
 }
@@ -556,12 +563,10 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 
 	jobID := takeFirst(orig.ID, uuid.New())
 	// Always set some tags to prevent Acquire from grabbing jobs it should not.
-	if !orig.StartedAt.Time.IsZero() {
-		if orig.Tags == nil {
-			orig.Tags = make(database.StringMap)
-		}
+	tags := takeFirstMap(orig.Tags, database.StringMap{"owner": "", "scope": "organization"})
+	if orig.Tags == nil && !orig.StartedAt.Time.IsZero() {
 		// Make sure when we acquire the job, we only get this one.
-		orig.Tags[jobID.String()] = "true"
+		tags[jobID.String()] = "true"
 	}
 
 	job, err := db.InsertProvisionerJob(genCtx, database.InsertProvisionerJobParams{
@@ -575,7 +580,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 		FileID:         takeFirst(orig.FileID, uuid.New()),
 		Type:           takeFirst(orig.Type, database.ProvisionerJobTypeWorkspaceBuild),
 		Input:          takeFirstSlice(orig.Input, []byte("{}")),
-		Tags:           orig.Tags,
+		Tags:           tags,
 		TraceMetadata:  pqtype.NullRawMessage{},
 	})
 	require.NoError(t, err, "insert job")
@@ -587,17 +592,18 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 		job, err = db.AcquireProvisionerJob(genCtx, database.AcquireProvisionerJobParams{
 			StartedAt:       orig.StartedAt,
 			OrganizationID:  job.OrganizationID,
-			Types:           []database.ProvisionerType{database.ProvisionerTypeEcho},
-			ProvisionerTags: must(json.Marshal(orig.Tags)),
-			WorkerID:        uuid.NullUUID{},
+			Types:           []database.ProvisionerType{job.Provisioner},
+			ProvisionerTags: must(json.Marshal(tags)),
+			WorkerID:        takeFirst(orig.WorkerID, uuid.NullUUID{}),
 		})
 		require.NoError(t, err)
 		// There is no easy way to make sure we acquire the correct job.
 		require.Equal(t, jobID, job.ID, "acquired incorrect job")
+		fmt.Printf("%#v\n", job)
 	}
 
 	if !orig.CompletedAt.Time.IsZero() || orig.Error.String != "" {
-		err := db.UpdateProvisionerJobWithCompleteByID(genCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
+		err = db.UpdateProvisionerJobWithCompleteByID(genCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID:          jobID,
 			UpdatedAt:   job.UpdatedAt,
 			CompletedAt: orig.CompletedAt,
@@ -607,7 +613,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 		require.NoError(t, err)
 	}
 	if !orig.CanceledAt.Time.IsZero() {
-		err := db.UpdateProvisionerJobWithCancelByID(genCtx, database.UpdateProvisionerJobWithCancelByIDParams{
+		err = db.UpdateProvisionerJobWithCancelByID(genCtx, database.UpdateProvisionerJobWithCancelByIDParams{
 			ID:          jobID,
 			CanceledAt:  orig.CanceledAt,
 			CompletedAt: orig.CompletedAt,
@@ -616,7 +622,7 @@ func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig data
 	}
 
 	job, err = db.GetProvisionerJobByID(genCtx, jobID)
-	require.NoError(t, err)
+	require.NoError(t, err, "get job: %s", jobID.String())
 
 	return job
 }
@@ -1108,6 +1114,12 @@ func takeFirstSlice[T any](values ...[]T) []T {
 	})
 }
 
+func takeFirstMap[T, E comparable](values ...map[T]E) map[T]E {
+	return takeFirstF(values, func(v map[T]E) bool {
+		return v != nil
+	})
+}
+
 // takeFirstF takes the first value that returns true
 func takeFirstF[Value any](values []Value, take func(v Value) bool) Value {
 	for _, v := range values {