@@ -1566,6 +1566,69 @@ func TestCompleteJob(t *testing.T) {
1566
1566
require .Contains (t ,job .Error .String ,`external auth provider "github" is not configured` )
1567
1567
})
1568
1568
1569
+ t .Run ("TemplateImport_WithAITasks" ,func (t * testing.T ) {
1570
+ t .Parallel ()
1571
+
1572
+ t .Run ("WithoutPromptParam" ,func (t * testing.T ) {
1573
+ t .Parallel ()
1574
+
1575
+ srv ,db ,_ ,pd := setup (t ,false ,& overrides {})
1576
+ jobID := uuid .New ()
1577
+ versionID := uuid .New ()
1578
+ err := db .InsertTemplateVersion (ctx , database.InsertTemplateVersionParams {
1579
+ ID :versionID ,
1580
+ JobID :jobID ,
1581
+ OrganizationID :pd .OrganizationID ,
1582
+ })
1583
+ require .NoError (t ,err )
1584
+ job ,err := db .InsertProvisionerJob (ctx , database.InsertProvisionerJobParams {
1585
+ ID :jobID ,
1586
+ Provisioner :database .ProvisionerTypeEcho ,
1587
+ Input : []byte (`{"template_version_id": "` + versionID .String ()+ `"}` ),
1588
+ StorageMethod :database .ProvisionerStorageMethodFile ,
1589
+ Type :database .ProvisionerJobTypeWorkspaceBuild ,
1590
+ OrganizationID :pd .OrganizationID ,
1591
+ Tags :pd .Tags ,
1592
+ })
1593
+ require .NoError (t ,err )
1594
+ _ ,err = db .AcquireProvisionerJob (ctx , database.AcquireProvisionerJobParams {
1595
+ OrganizationID :pd .OrganizationID ,
1596
+ WorkerID : uuid.NullUUID {
1597
+ UUID :pd .ID ,
1598
+ Valid :true ,
1599
+ },
1600
+ Types : []database.ProvisionerType {database .ProvisionerTypeEcho },
1601
+ StartedAt : sql.NullTime {
1602
+ Time :dbtime .Now (),
1603
+ Valid :true ,
1604
+ },
1605
+ ProvisionerTags :must (json .Marshal (job .Tags )),
1606
+ })
1607
+ require .NoError (t ,err )
1608
+ completeJob := func () {
1609
+ _ ,err = srv .CompleteJob (ctx ,& proto.CompletedJob {
1610
+ JobId :job .ID .String (),
1611
+ Type :& proto.CompletedJob_TemplateImport_ {
1612
+ TemplateImport :& proto.CompletedJob_TemplateImport {
1613
+ StartResources : []* sdkproto.Resource {{
1614
+ Name :"my_task" ,
1615
+ Type :"coder_ai_task" ,
1616
+ }},
1617
+ StopResources : []* sdkproto.Resource {},
1618
+ HasAiTasks :true ,
1619
+ Plan : []byte ("{}" ),
1620
+ },
1621
+ },
1622
+ })
1623
+ require .NoError (t ,err )
1624
+ }
1625
+ completeJob ()
1626
+ job ,err = db .GetProvisionerJobByID (ctx ,job .ID )
1627
+ require .NoError (t ,err )
1628
+ require .Contains (t ,job .Error .String ,`external auth provider "github" is not configured` )
1629
+ })
1630
+ })
1631
+
1569
1632
t .Run ("TemplateImport_WithGitAuth" ,func (t * testing.T ) {
1570
1633
t .Parallel ()
1571
1634
srv ,db ,_ ,pd := setup (t ,false ,& overrides {
@@ -2441,6 +2504,233 @@ func TestCompleteJob(t *testing.T) {
2441
2504
testutil .RequireReceive (ctx ,t ,done )
2442
2505
require .Equal (t ,replacements ,orchestrator .replacements )
2443
2506
})
2507
+
2508
+ t .Run ("AITasks" ,func (t * testing.T ) {
2509
+ t .Parallel ()
2510
+
2511
+ // has_ai_task has a default value of nil, but once the template import completes it will have a value;
2512
+ // it is set to "true" if the template has any coder_ai_task resources defined.
2513
+ t .Run ("TemplateImport" ,func (t * testing.T ) {
2514
+ type testcase struct {
2515
+ name string
2516
+ input * proto.CompletedJob_TemplateImport
2517
+ expected bool
2518
+ }
2519
+
2520
+ for _ ,tc := range []testcase {
2521
+ {
2522
+ name :"has_ai_task is false by default" ,
2523
+ input :& proto.CompletedJob_TemplateImport {
2524
+ // HasAiTasks is not set.
2525
+ Plan : []byte ("{}" ),
2526
+ },
2527
+ expected :false ,
2528
+ },
2529
+ {
2530
+ name :"has_ai_task gets set to true" ,
2531
+ input :& proto.CompletedJob_TemplateImport {
2532
+ HasAiTasks :true ,
2533
+ Plan : []byte ("{}" ),
2534
+ },
2535
+ expected :true ,
2536
+ },
2537
+ } {
2538
+ t .Run (tc .name ,func (t * testing.T ) {
2539
+ t .Parallel ()
2540
+
2541
+ srv ,db ,_ ,pd := setup (t ,false ,& overrides {})
2542
+
2543
+ importJobID := uuid .New ()
2544
+ tvID := uuid .New ()
2545
+ template := dbgen .Template (t ,db , database.Template {
2546
+ Name :"template" ,
2547
+ Provisioner :database .ProvisionerTypeEcho ,
2548
+ OrganizationID :pd .OrganizationID ,
2549
+ })
2550
+ version := dbgen .TemplateVersion (t ,db , database.TemplateVersion {
2551
+ ID :tvID ,
2552
+ OrganizationID :pd .OrganizationID ,
2553
+ TemplateID : uuid.NullUUID {
2554
+ UUID :template .ID ,
2555
+ Valid :true ,
2556
+ },
2557
+ JobID :importJobID ,
2558
+ })
2559
+ _ = version
2560
+
2561
+ ctx := testutil .Context (t ,testutil .WaitShort )
2562
+ job ,err := db .InsertProvisionerJob (ctx , database.InsertProvisionerJobParams {
2563
+ ID :importJobID ,
2564
+ CreatedAt :dbtime .Now (),
2565
+ UpdatedAt :dbtime .Now (),
2566
+ OrganizationID :pd .OrganizationID ,
2567
+ InitiatorID :uuid .New (),
2568
+ Input :must (json .Marshal (provisionerdserver.TemplateVersionImportJob {
2569
+ TemplateVersionID :tvID ,
2570
+ })),
2571
+ Provisioner :database .ProvisionerTypeEcho ,
2572
+ StorageMethod :database .ProvisionerStorageMethodFile ,
2573
+ Type :database .ProvisionerJobTypeTemplateVersionImport ,
2574
+ Tags :pd .Tags ,
2575
+ })
2576
+ require .NoError (t ,err )
2577
+
2578
+ _ ,err = db .AcquireProvisionerJob (ctx , database.AcquireProvisionerJobParams {
2579
+ OrganizationID :pd .OrganizationID ,
2580
+ WorkerID : uuid.NullUUID {
2581
+ UUID :pd .ID ,
2582
+ Valid :true ,
2583
+ },
2584
+ Types : []database.ProvisionerType {database .ProvisionerTypeEcho },
2585
+ ProvisionerTags :must (json .Marshal (job .Tags )),
2586
+ StartedAt : sql.NullTime {Time :job .CreatedAt ,Valid :true },
2587
+ })
2588
+ require .NoError (t ,err )
2589
+
2590
+ version ,err = db .GetTemplateVersionByID (ctx ,tvID )
2591
+ require .NoError (t ,err )
2592
+ require .False (t ,version .HasAITask .Valid )// Value should be nil (i.e. valid = false).
2593
+
2594
+ completedJob := proto.CompletedJob {
2595
+ JobId :job .ID .String (),
2596
+ Type :& proto.CompletedJob_TemplateImport_ {
2597
+ TemplateImport :tc .input ,
2598
+ },
2599
+ }
2600
+ _ ,err = srv .CompleteJob (ctx ,& completedJob )
2601
+ require .NoError (t ,err )
2602
+
2603
+ version ,err = db .GetTemplateVersionByID (ctx ,tvID )
2604
+ require .NoError (t ,err )
2605
+ require .True (t ,version .HasAITask .Valid )// We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true.
2606
+ require .Equal (t ,tc .expected ,version .HasAITask .Bool )
2607
+ })
2608
+ }
2609
+ })
2610
+
2611
+ // has_ai_task has a default value of nil, but once the workspace build completes it will have a value;
2612
+ // it is set to "true" if the related template has any coder_ai_task resources defined, and its sidebar app ID
2613
+ // will be set as well in that case.
2614
+ t .Run ("WorkspaceBuild" ,func (t * testing.T ) {
2615
+ type testcase struct {
2616
+ name string
2617
+ input * proto.CompletedJob_WorkspaceBuild
2618
+ expected bool
2619
+ }
2620
+
2621
+ sidebarAppID := uuid .NewString ()
2622
+ for _ ,tc := range []testcase {
2623
+ {
2624
+ name :"has_ai_task is false by default" ,
2625
+ input :& proto.CompletedJob_WorkspaceBuild {
2626
+ // No AiTasks defined.
2627
+ },
2628
+ expected :false ,
2629
+ },
2630
+ {
2631
+ name :"has_ai_task is set to true" ,
2632
+ input :& proto.CompletedJob_WorkspaceBuild {
2633
+ AiTasks : []* sdkproto.AITask {
2634
+ {
2635
+ Id :uuid .NewString (),
2636
+ SidebarApp :& sdkproto.AITaskSidebarApp {
2637
+ Id :sidebarAppID ,
2638
+ },
2639
+ },
2640
+ },
2641
+ },
2642
+ expected :true ,
2643
+ },
2644
+ } {
2645
+ t .Run (tc .name ,func (t * testing.T ) {
2646
+ t .Parallel ()
2647
+
2648
+ srv ,db ,_ ,pd := setup (t ,false ,& overrides {})
2649
+
2650
+ importJobID := uuid .New ()
2651
+ tvID := uuid .New ()
2652
+ template := dbgen .Template (t ,db , database.Template {
2653
+ Name :"template" ,
2654
+ Provisioner :database .ProvisionerTypeEcho ,
2655
+ OrganizationID :pd .OrganizationID ,
2656
+ })
2657
+ version := dbgen .TemplateVersion (t ,db , database.TemplateVersion {
2658
+ ID :tvID ,
2659
+ OrganizationID :pd .OrganizationID ,
2660
+ TemplateID : uuid.NullUUID {
2661
+ UUID :template .ID ,
2662
+ Valid :true ,
2663
+ },
2664
+ JobID :importJobID ,
2665
+ })
2666
+ user := dbgen .User (t ,db , database.User {})
2667
+ workspaceTable := dbgen .Workspace (t ,db , database.WorkspaceTable {
2668
+ TemplateID :template .ID ,
2669
+ OwnerID :user .ID ,
2670
+ OrganizationID :pd .OrganizationID ,
2671
+ })
2672
+ build := dbgen .WorkspaceBuild (t ,db , database.WorkspaceBuild {
2673
+ WorkspaceID :workspaceTable .ID ,
2674
+ TemplateVersionID :version .ID ,
2675
+ InitiatorID :user .ID ,
2676
+ Transition :database .WorkspaceTransitionStart ,
2677
+ })
2678
+
2679
+ ctx := testutil .Context (t ,testutil .WaitShort )
2680
+ job ,err := db .InsertProvisionerJob (ctx , database.InsertProvisionerJobParams {
2681
+ ID :importJobID ,
2682
+ CreatedAt :dbtime .Now (),
2683
+ UpdatedAt :dbtime .Now (),
2684
+ OrganizationID :pd .OrganizationID ,
2685
+ InitiatorID :uuid .New (),
2686
+ Input :must (json .Marshal (provisionerdserver.WorkspaceProvisionJob {
2687
+ WorkspaceBuildID :build .ID ,
2688
+ LogLevel :"DEBUG" ,
2689
+ })),
2690
+ Provisioner :database .ProvisionerTypeEcho ,
2691
+ StorageMethod :database .ProvisionerStorageMethodFile ,
2692
+ Type :database .ProvisionerJobTypeWorkspaceBuild ,
2693
+ Tags :pd .Tags ,
2694
+ })
2695
+ require .NoError (t ,err )
2696
+
2697
+ _ ,err = db .AcquireProvisionerJob (ctx , database.AcquireProvisionerJobParams {
2698
+ OrganizationID :pd .OrganizationID ,
2699
+ WorkerID : uuid.NullUUID {
2700
+ UUID :pd .ID ,
2701
+ Valid :true ,
2702
+ },
2703
+ Types : []database.ProvisionerType {database .ProvisionerTypeEcho },
2704
+ ProvisionerTags :must (json .Marshal (job .Tags )),
2705
+ StartedAt : sql.NullTime {Time :job .CreatedAt ,Valid :true },
2706
+ })
2707
+ require .NoError (t ,err )
2708
+
2709
+ build ,err = db .GetWorkspaceBuildByID (ctx ,build .ID )
2710
+ require .NoError (t ,err )
2711
+ require .False (t ,build .HasAITask .Valid )// Value should be nil (i.e. valid = false).
2712
+
2713
+ completedJob := proto.CompletedJob {
2714
+ JobId :job .ID .String (),
2715
+ Type :& proto.CompletedJob_WorkspaceBuild_ {
2716
+ WorkspaceBuild :tc .input ,
2717
+ },
2718
+ }
2719
+ _ ,err = srv .CompleteJob (ctx ,& completedJob )
2720
+ require .NoError (t ,err )
2721
+
2722
+ build ,err = db .GetWorkspaceBuildByID (ctx ,build .ID )
2723
+ require .NoError (t ,err )
2724
+ require .True (t ,build .HasAITask .Valid )// We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true.
2725
+ require .Equal (t ,tc .expected ,build .HasAITask .Bool )
2726
+
2727
+ if tc .expected {
2728
+ require .Equal (t ,sidebarAppID ,build .AITaskSidebarAppID .UUID .String ())
2729
+ }
2730
+ })
2731
+ }
2732
+ })
2733
+ })
2444
2734
}
2445
2735
2446
2736
type mockPrebuildsOrchestrator struct {