@@ -12,7 +12,7 @@ import (
12
12
13
13
"cdr.dev/slog/sloggers/slogtest"
14
14
"github.com/coder/coder/coderd/database"
15
- "github.com/coder/coder/coderd/database/dbtestutil "
15
+ "github.com/coder/coder/coderd/database/databasefake "
16
16
"github.com/coder/coder/coderd/provisionerdserver"
17
17
"github.com/coder/coder/coderd/telemetry"
18
18
"github.com/coder/coder/codersdk"
@@ -498,9 +498,268 @@ func TestFailJob(t *testing.T) {
498
498
})
499
499
}
500
500
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
+
501
759
func setup (t * testing.T )* provisionerdserver.Server {
502
760
t .Helper ()
503
- db ,pubsub := dbtestutil .NewDB (t )
761
+ db := databasefake .New ()
762
+ pubsub := database .NewPubsubInMemory ()
504
763
505
764
return & provisionerdserver.Server {
506
765
ID :uuid .New (),