add tests

coder · SasSwart · Dec 11, 2024 · Nov 29, 2024 · Dec 3, 2024 · Dec 4, 2024
commit 58851972586ada4b66f33f314d56e6254ce0b36d
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -503,6 +503,46 @@ func GroupMember(t testing.TB, db database.Store, member database.GroupMemberTab
 	return groupMember
 }
 
+// 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 {
+	t.Helper()
+
+	if daemon.KeyID == uuid.Nil {
+		key, err := db.InsertProvisionerKey(genCtx, database.InsertProvisionerKeyParams{
+			ID:             uuid.New(),
+			Name:           daemon.Name + "-key",
+			OrganizationID: daemon.OrganizationID,
+			HashedSecret:   []byte("secret"),
+			CreatedAt:      dbtime.Now(),
+			Tags:           daemon.Tags,
+		})
+		require.NoError(t, err)
+		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,
+	})
+	require.NoError(t, err)
+	return d
+}
+
 // ProvisionerJob is a bit more involved to get the values such as "completedAt", "startedAt", "cancelledAt" set.  ps
 // can be set to nil if you are SURE that you don't require a provisionerdaemon to acquire the job in your test.
 func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig database.ProvisionerJob) database.ProvisionerJob {

diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -27,6 +27,7 @@ import (
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/rbac/policy"
+	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/testutil"
 )
 
@@ -211,6 +212,145 @@ func TestGetDeploymentWorkspaceAgentUsageStats(t *testing.T) {
 	})
 }
 
+func TestGetEligibleProvisionerDaemonsByProvisionerJobIDs(t *testing.T) {
+	t.Parallel()
+
+	t.Run("NoJobsReturnsEmpty", func(t *testing.T) {
+		t.Parallel()
+		db, _ := dbtestutil.NewDB(t)
+		daemons, err := db.GetEligibleProvisionerDaemonsByProvisionerJobIDs(context.Background(), []uuid.UUID{})
+		require.NoError(t, err)
+		require.Empty(t, daemons)
+	})
+
+	t.Run("MatchesProvisionerType", func(t *testing.T) {
+		t.Parallel()
+		db, _ := dbtestutil.NewDB(t)
+		org := dbgen.Organization(t, db, database.Organization{})
+
+		job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			OrganizationID: org.ID,
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			Provisioner:    database.ProvisionerTypeEcho,
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		matchingDaemon := dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "matching-daemon",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "non-matching-daemon",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeTerraform},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		daemons, err := db.GetEligibleProvisionerDaemonsByProvisionerJobIDs(context.Background(), []uuid.UUID{job.ID})
+		require.NoError(t, err)
+		require.Len(t, daemons, 1)
+		require.Equal(t, matchingDaemon.ID, daemons[0].ProvisionerDaemon.ID)
+	})
+
+	t.Run("MatchesOrganizationScope", func(t *testing.T) {
+		t.Parallel()
+		db, _ := dbtestutil.NewDB(t)
+		org := dbgen.Organization(t, db, database.Organization{})
+
+		job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			OrganizationID: org.ID,
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			Provisioner:    database.ProvisionerTypeEcho,
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+				provisionersdk.TagOwner: "",
+			},
+		})
+
+		orgDaemon := dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "org-daemon",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+				provisionersdk.TagOwner: "",
+			},
+		})
+
+		dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "user-daemon",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeUser,
+			},
+		})
+
+		daemons, err := db.GetEligibleProvisionerDaemonsByProvisionerJobIDs(context.Background(), []uuid.UUID{job.ID})
+		require.NoError(t, err)
+		require.Len(t, daemons, 1)
+		require.Equal(t, orgDaemon.ID, daemons[0].ProvisionerDaemon.ID)
+	})
+
+	t.Run("MatchesMultipleProvisioners", func(t *testing.T) {
+		t.Parallel()
+		db, _ := dbtestutil.NewDB(t)
+		org := dbgen.Organization(t, db, database.Organization{})
+
+		job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			OrganizationID: org.ID,
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			Provisioner:    database.ProvisionerTypeEcho,
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		daemon1 := dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "daemon-1",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		daemon2 := dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "daemon-2",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeEcho},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		dbgen.ProvisionerDaemon(t, db, database.ProvisionerDaemon{
+			Name:           "daemon-3",
+			OrganizationID: org.ID,
+			Provisioners:   []database.ProvisionerType{database.ProvisionerTypeTerraform},
+			Tags: database.StringMap{
+				provisionersdk.TagScope: provisionersdk.ScopeOrganization,
+			},
+		})
+
+		daemons, err := db.GetEligibleProvisionerDaemonsByProvisionerJobIDs(context.Background(), []uuid.UUID{job.ID})
+		require.NoError(t, err)
+		require.Len(t, daemons, 2)
+
+		daemonIDs := []uuid.UUID{daemons[0].ProvisionerDaemon.ID, daemons[1].ProvisionerDaemon.ID}
+		require.ElementsMatch(t, []uuid.UUID{daemon1.ID, daemon2.ID}, daemonIDs)
+	})
+}
+
 func TestGetWorkspaceAgentUsageStats(t *testing.T) {
 	t.Parallel()
 

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/provisionerdaemons.sql b/coderd/database/queries/provisionerdaemons.sql
@@ -24,6 +24,7 @@ FROM
 JOIN
     provisioner_daemons ON provisioner_daemons.organization_id = provisioner_jobs.organization_id
     AND provisioner_tagset_contains(provisioner_daemons.tags::tagset, provisioner_jobs.tags::tagset)
+    AND provisioner_jobs.provisioner = ANY(provisioner_daemons.provisioners)
 WHERE
     provisioner_jobs.id = ANY(@provisioner_job_ids :: uuid[]);