Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a57a934

Browse files
kylecarbskylecarbs
committed
Add test for workspace resources
1 parent dfa73b3 commit a57a934

File tree

2 files changed

+269
-7
lines changed

2 files changed

+269
-7
lines changed

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -489,7 +489,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
489489
return nil, xerrors.Errorf("get job by id: %w", err)
490490
}
491491
if job.WorkerID.UUID.String() != server.ID.String() {
492-
return nil, xerrors.Errorf("you don't have permission to update this job")
492+
return nil, xerrors.Errorf("you don't own this job")
493493
}
494494

495495
telemetrySnapshot := &telemetry.Snapshot{}
@@ -509,7 +509,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
509509
slog.F("resource_type", resource.Type),
510510
slog.F("transition", transition))
511511

512-
err = insertWorkspaceResource(ctx, server.Database, jobID, transition, resource, telemetrySnapshot)
512+
err = InsertWorkspaceResource(ctx, server.Database, jobID, transition, resource, telemetrySnapshot)
513513
if err != nil {
514514
return nil, xerrors.Errorf("insert resource: %w", err)
515515
}
@@ -578,7 +578,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
578578
}
579579
// This could be a bulk insert to improve performance.
580580
for _, protoResource := range jobType.WorkspaceBuild.Resources {
581-
err = insertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
581+
err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
582582
if err != nil {
583583
return xerrors.Errorf("insert provisioner job: %w", err)
584584
}
@@ -614,7 +614,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
614614
slog.F("resource_name", resource.Name),
615615
slog.F("resource_type", resource.Type))
616616

617-
err = insertWorkspaceResource(ctx, server.Database, jobID, database.WorkspaceTransitionStart, resource, telemetrySnapshot)
617+
err = InsertWorkspaceResource(ctx, server.Database, jobID, database.WorkspaceTransitionStart, resource, telemetrySnapshot)
618618
if err != nil {
619619
return nil, xerrors.Errorf("insert resource: %w", err)
620620
}
@@ -637,6 +637,9 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
637637
}
638638

639639
default:
640+
if completed.Type == nil {
641+
return nil, xerrors.Errorf("type payload must be provided")
642+
}
640643
return nil, xerrors.Errorf("unknown job type %q; ensure coderd and provisionerd versions match",
641644
reflect.TypeOf(completed.Type).String())
642645
}
@@ -655,7 +658,7 @@ func (server *Server) CompleteJob(ctx context.Context, completed *proto.Complete
655658
return &proto.Empty{}, nil
656659
}
657660

658-
func insertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error {
661+
func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error {
659662
resource, err := db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
660663
ID: uuid.New(),
661664
CreatedAt: database.Now(),

‎coderd/provisionerdserver/provisionerdserver_test.go

Lines changed: 261 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212

1313
"cdr.dev/slog/sloggers/slogtest"
1414
"github.com/coder/coder/coderd/database"
15-
"github.com/coder/coder/coderd/database/dbtestutil"
15+
"github.com/coder/coder/coderd/database/databasefake"
1616
"github.com/coder/coder/coderd/provisionerdserver"
1717
"github.com/coder/coder/coderd/telemetry"
1818
"github.com/coder/coder/codersdk"
@@ -498,9 +498,268 @@ func TestFailJob(t *testing.T) {
498498
})
499499
}
500500

501+
func TestCompleteJob(t *testing.T) {
502+
t.Parallel()
503+
ctx := context.Background()
504+
t.Run("NotFound", func(t *testing.T) {
505+
t.Parallel()
506+
srv := setup(t)
507+
_, err := srv.CompleteJob(ctx, &proto.CompletedJob{
508+
JobId: "hello",
509+
})
510+
require.ErrorContains(t, err, "invalid UUID")
511+
512+
_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
513+
JobId: uuid.NewString(),
514+
})
515+
require.ErrorContains(t, err, "no rows in result set")
516+
})
517+
// This test prevents runners from updating jobs they don't own!
518+
t.Run("NotOwner", func(t *testing.T) {
519+
t.Parallel()
520+
srv := setup(t)
521+
job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
522+
ID: uuid.New(),
523+
Provisioner: database.ProvisionerTypeEcho,
524+
})
525+
require.NoError(t, err)
526+
_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
527+
WorkerID: uuid.NullUUID{
528+
UUID: uuid.New(),
529+
Valid: true,
530+
},
531+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
532+
})
533+
require.NoError(t, err)
534+
_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
535+
JobId: job.ID.String(),
536+
})
537+
require.ErrorContains(t, err, "you don't own this job")
538+
})
539+
t.Run("TemplateImport", func(t *testing.T) {
540+
t.Parallel()
541+
srv := setup(t)
542+
job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
543+
ID: uuid.New(),
544+
Provisioner: database.ProvisionerTypeEcho,
545+
})
546+
require.NoError(t, err)
547+
_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
548+
WorkerID: uuid.NullUUID{
549+
UUID: srv.ID,
550+
Valid: true,
551+
},
552+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
553+
})
554+
require.NoError(t, err)
555+
_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
556+
JobId: job.ID.String(),
557+
Type: &proto.CompletedJob_TemplateImport_{
558+
TemplateImport: &proto.CompletedJob_TemplateImport{
559+
StartResources: []*sdkproto.Resource{{
560+
Name: "hello",
561+
Type: "aws_instance",
562+
}},
563+
StopResources: []*sdkproto.Resource{},
564+
},
565+
},
566+
})
567+
require.NoError(t, err)
568+
})
569+
t.Run("WorkspaceBuild", func(t *testing.T) {
570+
t.Parallel()
571+
srv := setup(t)
572+
workspace, err := srv.Database.InsertWorkspace(ctx, database.InsertWorkspaceParams{
573+
ID: uuid.New(),
574+
})
575+
require.NoError(t, err)
576+
build, err := srv.Database.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
577+
ID: uuid.New(),
578+
WorkspaceID: workspace.ID,
579+
Transition: database.WorkspaceTransitionDelete,
580+
})
581+
require.NoError(t, err)
582+
input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
583+
WorkspaceBuildID: build.ID,
584+
})
585+
require.NoError(t, err)
586+
job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
587+
ID: uuid.New(),
588+
Provisioner: database.ProvisionerTypeEcho,
589+
Input: input,
590+
})
591+
require.NoError(t, err)
592+
_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
593+
WorkerID: uuid.NullUUID{
594+
UUID: srv.ID,
595+
Valid: true,
596+
},
597+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
598+
})
599+
require.NoError(t, err)
600+
601+
publishedWorkspace := make(chan struct{})
602+
closeWorkspaceSubscribe, err := srv.Pubsub.Subscribe(codersdk.WorkspaceNotifyChannel(build.WorkspaceID), func(_ context.Context, _ []byte) {
603+
close(publishedWorkspace)
604+
})
605+
require.NoError(t, err)
606+
defer closeWorkspaceSubscribe()
607+
publishedLogs := make(chan struct{})
608+
closeLogsSubscribe, err := srv.Pubsub.Subscribe(provisionerdserver.ProvisionerJobLogsNotifyChannel(job.ID), func(_ context.Context, _ []byte) {
609+
close(publishedLogs)
610+
})
611+
require.NoError(t, err)
612+
defer closeLogsSubscribe()
613+
614+
_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
615+
JobId: job.ID.String(),
616+
Type: &proto.CompletedJob_WorkspaceBuild_{
617+
WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{
618+
State: []byte{},
619+
Resources: []*sdkproto.Resource{{
620+
Name: "example",
621+
Type: "aws_instance",
622+
}},
623+
},
624+
},
625+
})
626+
require.NoError(t, err)
627+
628+
<-publishedWorkspace
629+
<-publishedLogs
630+
631+
workspace, err = srv.Database.GetWorkspaceByID(ctx, workspace.ID)
632+
require.NoError(t, err)
633+
require.True(t, workspace.Deleted)
634+
})
635+
636+
t.Run("TemplateDryRun", func(t *testing.T) {
637+
t.Parallel()
638+
srv := setup(t)
639+
job, err := srv.Database.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
640+
ID: uuid.New(),
641+
Provisioner: database.ProvisionerTypeEcho,
642+
})
643+
require.NoError(t, err)
644+
_, err = srv.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
645+
WorkerID: uuid.NullUUID{
646+
UUID: srv.ID,
647+
Valid: true,
648+
},
649+
Types: []database.ProvisionerType{database.ProvisionerTypeEcho},
650+
})
651+
require.NoError(t, err)
652+
653+
_, err = srv.CompleteJob(ctx, &proto.CompletedJob{
654+
JobId: job.ID.String(),
655+
Type: &proto.CompletedJob_TemplateDryRun_{
656+
TemplateDryRun: &proto.CompletedJob_TemplateDryRun{
657+
Resources: []*sdkproto.Resource{{
658+
Name: "something",
659+
Type: "aws_instance",
660+
}},
661+
},
662+
},
663+
})
664+
require.NoError(t, err)
665+
})
666+
}
667+
668+
func TestInsertWorkspaceResource(t *testing.T) {
669+
t.Parallel()
670+
ctx := context.Background()
671+
insert := func(db database.Store, jobID uuid.UUID, resource *sdkproto.Resource) error {
672+
return provisionerdserver.InsertWorkspaceResource(ctx, db, jobID, database.WorkspaceTransitionStart, resource, &telemetry.Snapshot{})
673+
}
674+
t.Run("NoAgents", func(t *testing.T) {
675+
t.Parallel()
676+
db := databasefake.New()
677+
job := uuid.New()
678+
err := insert(db, job, &sdkproto.Resource{
679+
Name: "something",
680+
Type: "aws_instance",
681+
})
682+
require.NoError(t, err)
683+
resources, err := db.GetWorkspaceResourcesByJobID(ctx, job)
684+
require.NoError(t, err)
685+
require.Len(t, resources, 1)
686+
})
687+
t.Run("InvalidAgentToken", func(t *testing.T) {
688+
t.Parallel()
689+
err := insert(databasefake.New(), uuid.New(), &sdkproto.Resource{
690+
Name: "something",
691+
Type: "aws_instance",
692+
Agents: []*sdkproto.Agent{{
693+
Auth: &sdkproto.Agent_Token{
694+
Token: "bananas",
695+
},
696+
}},
697+
})
698+
require.ErrorContains(t, err, "invalid UUID length")
699+
})
700+
t.Run("DuplicateApps", func(t *testing.T) {
701+
t.Parallel()
702+
err := insert(databasefake.New(), uuid.New(), &sdkproto.Resource{
703+
Name: "something",
704+
Type: "aws_instance",
705+
Agents: []*sdkproto.Agent{{
706+
Apps: []*sdkproto.App{{
707+
Slug: "a",
708+
}, {
709+
Slug: "a",
710+
}},
711+
}},
712+
})
713+
require.ErrorContains(t, err, "duplicate app slug")
714+
})
715+
t.Run("Success", func(t *testing.T) {
716+
t.Parallel()
717+
db := databasefake.New()
718+
job := uuid.New()
719+
err := insert(db, job, &sdkproto.Resource{
720+
Name: "something",
721+
Type: "aws_instance",
722+
Agents: []*sdkproto.Agent{{
723+
Name: "dev",
724+
Env: map[string]string{
725+
"something": "test",
726+
},
727+
StartupScript: "value",
728+
OperatingSystem: "linux",
729+
Architecture: "amd64",
730+
Auth: &sdkproto.Agent_Token{
731+
Token: uuid.NewString(),
732+
},
733+
Apps: []*sdkproto.App{{
734+
Slug: "a",
735+
}},
736+
}},
737+
})
738+
require.NoError(t, err)
739+
resources, err := db.GetWorkspaceResourcesByJobID(ctx, job)
740+
require.NoError(t, err)
741+
require.Len(t, resources, 1)
742+
agents, err := db.GetWorkspaceAgentsByResourceIDs(ctx, []uuid.UUID{resources[0].ID})
743+
require.NoError(t, err)
744+
require.Len(t, agents, 1)
745+
agent := agents[0]
746+
require.Equal(t, "amd64", agent.Architecture)
747+
require.Equal(t, "linux", agent.OperatingSystem)
748+
require.Equal(t, "value", agent.StartupScript.String)
749+
want, err := json.Marshal(map[string]string{
750+
"something": "test",
751+
})
752+
require.NoError(t, err)
753+
got, err := agent.EnvironmentVariables.RawMessage.MarshalJSON()
754+
require.NoError(t, err)
755+
require.Equal(t, want, got)
756+
})
757+
}
758+
501759
func setup(t *testing.T) *provisionerdserver.Server {
502760
t.Helper()
503-
db, pubsub := dbtestutil.NewDB(t)
761+
db := databasefake.New()
762+
pubsub := database.NewPubsubInMemory()
504763

505764
return &provisionerdserver.Server{
506765
ID: uuid.New(),

0 commit comments

Comments
 (0)