Add test for workspace resources

coder · kylecarbs · Nov 8, 2022 · Nov 7, 2022 · Nov 7, 2022 · Nov 7, 2022
commit a57a934034616a7516ece1eca0aa09df2fae395c
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -489,7 +489,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 		return nil, xerrors.Errorf("get job by id: %w", err)
 	}
 	if job.WorkerID.UUID.String() != server.ID.String() {
-		return nil, xerrors.Errorf("you don't have permission to update this job")
+		return nil, xerrors.Errorf("you don't own this job")
 	}
 
 	telemetrySnapshot := &telemetry.Snapshot{}
@@ -509,7 +509,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 					slog.F("resource_type", resource.Type),
 					slog.F("transition", transition))
 
-				err = insertWorkspaceResource(ctx, server.Database, jobID, transition, resource, telemetrySnapshot)
+				err = InsertWorkspaceResource(ctx, server.Database, jobID, transition, resource, telemetrySnapshot)
 				if err != nil {
 					return nil, xerrors.Errorf("insert resource: %w", err)
 				}
@@ -578,7 +578,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 			}
 			// This could be a bulk insert to improve performance.
 			for _, protoResource := range jobType.WorkspaceBuild.Resources {
-				err = insertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
+				err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
 				if err != nil {
 					return xerrors.Errorf("insert provisioner job: %w", err)
 				}
@@ -614,7 +614,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 				slog.F("resource_name", resource.Name),
 				slog.F("resource_type", resource.Type))
 
-			err = insertWorkspaceResource(ctx, server.Database, jobID, database.WorkspaceTransitionStart, resource, telemetrySnapshot)
+			err = InsertWorkspaceResource(ctx, server.Database, jobID, database.WorkspaceTransitionStart, resource, telemetrySnapshot)
 			if err != nil {
 				return nil, xerrors.Errorf("insert resource: %w", err)
 			}
@@ -637,6 +637,9 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 		}
 
 	default:
+		if completed.Type == nil {
+			return nil, xerrors.Errorf("type payload must be provided")
+		}
 		return nil, xerrors.Errorf("unknown job type %q; ensure coderd and provisionerd versions match",
 			reflect.TypeOf(completed.Type).String())
 	}
@@ -655,7 +658,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
 	return &proto.Empty{}, nil
 }
 
-func insertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error {
+func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error {
 	resource, err := db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
 		ID:         uuid.New(),
 		CreatedAt:  database.Now(),

diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -12,7 +12,7 @@ import (
 
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/coderd/database"
-	"github.com/coder/coder/coderd/database/dbtestutil"
+	"github.com/coder/coder/coderd/database/databasefake"
 	"github.com/coder/coder/coderd/provisionerdserver"
 	"github.com/coder/coder/coderd/telemetry"
 	"github.com/coder/coder/codersdk"
@@ -498,9 +498,268 @@ func TestFailJob(t *testing.T) {
 	})
 }
 
+func TestCompleteJob(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	t.Run("NotFound", func(t *testing.T) {
+		t.Parallel()
+		srv := setup(t)
+		_, err := srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: "hello",
+		})
+		require.ErrorContains(t, err, "invalid UUID")
+
+		_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: uuid.NewString(),
+		})
+		require.ErrorContains(t, err, "no rows in result set")
+	})
+	// This test prevents runners from updating jobs they don't own!
+	t.Run("NotOwner", func(t *testing.T) {
+		t.Parallel()
+		srv := setup(t)
+		job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+			ID:          uuid.New(),
+			Provisioner: database.ProvisionerTypeEcho,
+		})
+		require.NoError(t, err)
+		_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
+			WorkerID: uuid.NullUUID{
+				UUID:  uuid.New(),
+				Valid: true,
+			},
+			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+		})
+		require.NoError(t, err)
+		_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: job.ID.String(),
+		})
+		require.ErrorContains(t, err, "you don't own this job")
+	})
+	t.Run("TemplateImport", func(t *testing.T) {
+		t.Parallel()
+		srv := setup(t)
+		job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+			ID:          uuid.New(),
+			Provisioner: database.ProvisionerTypeEcho,
+		})
+		require.NoError(t, err)
+		_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
+			WorkerID: uuid.NullUUID{
+				UUID:  srv.ID,
+				Valid: true,
+			},
+			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+		})
+		require.NoError(t, err)
+		_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: job.ID.String(),
+			Type: &proto.CompletedJob_TemplateImport_{
+				TemplateImport: &proto.CompletedJob_TemplateImport{
+					StartResources: []*sdkproto.Resource{{
+						Name: "hello",
+						Type: "aws_instance",
+					}},
+					StopResources: []*sdkproto.Resource{},
+				},
+			},
+		})
+		require.NoError(t, err)
+	})
+	t.Run("WorkspaceBuild", func(t *testing.T) {
+		t.Parallel()
+		srv := setup(t)
+		workspace, err := srv.Database.InsertWorkspace(ctx, database.InsertWorkspaceParams{
+			ID: uuid.New(),
+		})
+		require.NoError(t, err)
+		build, err := srv.Database.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
+			ID:          uuid.New(),
+			WorkspaceID: workspace.ID,
+			Transition:  database.WorkspaceTransitionDelete,
+		})
+		require.NoError(t, err)
+		input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
+			WorkspaceBuildID: build.ID,
+		})
+		require.NoError(t, err)
+		job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+			ID:          uuid.New(),
+			Provisioner: database.ProvisionerTypeEcho,
+			Input:       input,
+		})
+		require.NoError(t, err)
+		_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
+			WorkerID: uuid.NullUUID{
+				UUID:  srv.ID,
+				Valid: true,
+			},
+			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+		})
+		require.NoError(t, err)
+
+		publishedWorkspace := make(chan struct{})
+		closeWorkspaceSubscribe, err := srv.Pubsub.Subscribe(codersdk.WorkspaceNotifyChannel(build.WorkspaceID), func(_ context.Context, _ []byte) {
+			close(publishedWorkspace)
+		})
+		require.NoError(t, err)
+		defer closeWorkspaceSubscribe()
+		publishedLogs := make(chan struct{})
+		closeLogsSubscribe, err := srv.Pubsub.Subscribe(provisionerdserver.ProvisionerJobLogsNotifyChannel(job.ID), func(_ context.Context, _ []byte) {
+			close(publishedLogs)
+		})
+		require.NoError(t, err)
+		defer closeLogsSubscribe()
+
+		_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: job.ID.String(),
+			Type: &proto.CompletedJob_WorkspaceBuild_{
+				WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{
+					State: []byte{},
+					Resources: []*sdkproto.Resource{{
+						Name: "example",
+						Type: "aws_instance",
+					}},
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		<-publishedWorkspace
+		<-publishedLogs
+
+		workspace, err = srv.Database.GetWorkspaceByID(ctx, workspace.ID)
+		require.NoError(t, err)
+		require.True(t, workspace.Deleted)
+	})
+
+	t.Run("TemplateDryRun", func(t *testing.T) {
+		t.Parallel()
+		srv := setup(t)
+		job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+			ID:          uuid.New(),
+			Provisioner: database.ProvisionerTypeEcho,
+		})
+		require.NoError(t, err)
+		_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
+			WorkerID: uuid.NullUUID{
+				UUID:  srv.ID,
+				Valid: true,
+			},
+			Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
+		})
+		require.NoError(t, err)
+
+		_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
+			JobId: job.ID.String(),
+			Type: &proto.CompletedJob_TemplateDryRun_{
+				TemplateDryRun: &proto.CompletedJob_TemplateDryRun{
+					Resources: []*sdkproto.Resource{{
+						Name: "something",
+						Type: "aws_instance",
+					}},
+				},
+			},
+		})
+		require.NoError(t, err)
+	})
+}
+
+func TestInsertWorkspaceResource(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	insert := func(db database.Store, jobID uuid.UUID, resource *sdkproto.Resource) error {
+		return provisionerdserver.InsertWorkspaceResource(ctx, db, jobID, database.WorkspaceTransitionStart, resource, &telemetry.Snapshot{})
+	}
+	t.Run("NoAgents", func(t *testing.T) {
+		t.Parallel()
+		db := databasefake.New()
+		job := uuid.New()
+		err := insert(db, job, &sdkproto.Resource{
+			Name: "something",
+			Type: "aws_instance",
+		})
+		require.NoError(t, err)
+		resources, err := db.GetWorkspaceResourcesByJobID(ctx, job)
+		require.NoError(t, err)
+		require.Len(t, resources, 1)
+	})
+	t.Run("InvalidAgentToken", func(t *testing.T) {
+		t.Parallel()
+		err := insert(databasefake.New(), uuid.New(), &sdkproto.Resource{
+			Name: "something",
+			Type: "aws_instance",
+			Agents: []*sdkproto.Agent{{
+				Auth: &sdkproto.Agent_Token{
+					Token: "bananas",
+				},
+			}},
+		})
+		require.ErrorContains(t, err, "invalid UUID length")
+	})
+	t.Run("DuplicateApps", func(t *testing.T) {
+		t.Parallel()
+		err := insert(databasefake.New(), uuid.New(), &sdkproto.Resource{
+			Name: "something",
+			Type: "aws_instance",
+			Agents: []*sdkproto.Agent{{
+				Apps: []*sdkproto.App{{
+					Slug: "a",
+				}, {
+					Slug: "a",
+				}},
+			}},
+		})
+		require.ErrorContains(t, err, "duplicate app slug")
+	})
+	t.Run("Success", func(t *testing.T) {
+		t.Parallel()
+		db := databasefake.New()
+		job := uuid.New()
+		err := insert(db, job, &sdkproto.Resource{
+			Name: "something",
+			Type: "aws_instance",
+			Agents: []*sdkproto.Agent{{
+				Name: "dev",
+				Env: map[string]string{
+					"something": "test",
+				},
+				StartupScript:   "value",
+				OperatingSystem: "linux",
+				Architecture:    "amd64",
+				Auth: &sdkproto.Agent_Token{
+					Token: uuid.NewString(),
+				},
+				Apps: []*sdkproto.App{{
+					Slug: "a",
+				}},
+			}},
+		})
+		require.NoError(t, err)
+		resources, err := db.GetWorkspaceResourcesByJobID(ctx, job)
+		require.NoError(t, err)
+		require.Len(t, resources, 1)
+		agents, err := db.GetWorkspaceAgentsByResourceIDs(ctx, []uuid.UUID{resources[0].ID})
+		require.NoError(t, err)
+		require.Len(t, agents, 1)
+		agent := agents[0]
+		require.Equal(t, "amd64", agent.Architecture)
+		require.Equal(t, "linux", agent.OperatingSystem)
+		require.Equal(t, "value", agent.StartupScript.String)
+		want, err := json.Marshal(map[string]string{
+			"something": "test",
+		})
+		require.NoError(t, err)
+		got, err := agent.EnvironmentVariables.RawMessage.MarshalJSON()
+		require.NoError(t, err)
+		require.Equal(t, want, got)
+	})
+}
+
 func setup(t *testing.T) *provisionerdserver.Server {
 	t.Helper()
-	db, pubsub := dbtestutil.NewDB(t)
+	db := databasefake.New()
+	pubsub := database.NewPubsubInMemory()
 
 	return &provisionerdserver.Server{
 		ID:           uuid.New(),