Update Duo dev rake task

* Now it can enable either the Duo Pro or Duo Enterprise add-on

doc/development/ai_features/index.md

@@ -68,6 +68,16 @@ In SaaS mode, membership to a group with Duo features enabled is what enables
 many AI features. Make sure that your test user is a member of the group with
 Duo features enabled (`test-group-name`).
 
+This Rake task creates Duo Enterprise add-on attached to that group.
+
+In case you need Duo Pro add-on attached, please use:
+
+```shell
+GITLAB_SIMULATE_SAAS=1 bundle exec 'rake gitlab:duo:setup[test-group-name,duo_pro]'
+```
+
+Duo Pro add-on serves smaller scope of features. Usage of add-on depends on what features you want to use.
+
 #### Option B: in Self-managed Mode
 
 **Why:** If you want to test something specific to self-managed, such as Custom
@@ -78,13 +88,23 @@ Models.
 Run the Rake task to set up Duo features for the instance:
 
 ```shell
-GITLAB_SIMULATE_SAAS=0 bundle exec 'rake gitlab:duo:setup'
+GITLAB_SIMULATE_SAAS=0 bundle exec 'rake gitlab:duo:setup_instance'
 ```
 
 ```shell
 gdk restart
 ```
 
+This Rake task creates Duo Enterprise add-on attached to your instance.
+
+In case you need Duo Pro add-on attached, please use:
+
+```shell
+GITLAB_SIMULATE_SAAS=0 bundle exec 'rake gitlab:duo:setup_instance[duo_pro]'
+```
+
+Duo Pro add-on serves smaller scope of features. Usage of add-on depends on what features you want to use.
+
 ### Recommended: Set `CLOUD_CONNECTOR_SELF_SIGN_TOKENS` environment variable
 
 **Why:** Setting this environment variable will allow the local GitLab instance to

ee/lib/gitlab/duo/developments/setup.rb

@@ -188,14 +188,16 @@ def ensure_license!
         def create_add_on_purchases!
           group = Group.find_by_full_path(@namespace) # will be nil for self-managed mode
 
-          # rubocop: disable Cop/DestroyAll -- For dev
-          ::GitlabSubscriptions::AddOnPurchase.destroy_all
-          # rubocop: enable Cop/DestroyAll
-          create_code_suggestions_purchase!(group)
-          create_enterprise_purchase!(group)
+          ::GitlabSubscriptions::AddOnPurchase.by_namespace(group).delete_all
+
+          if args[:add_on] == 'duo_pro'
+            create_duo_pro_purchase!(group)
+          else
+            create_enterprise_purchase!(group)
+          end
         end
 
-        def create_code_suggestions_purchase!(group)
+        def create_duo_pro_purchase!(group)
           add_on = ::GitlabSubscriptions::AddOn.find_or_create_by_name(:code_suggestions)
 
           response = ::GitlabSubscriptions::AddOnPurchases::CreateService.new(group, add_on, {
@@ -206,6 +208,8 @@ def create_code_suggestions_purchase!(group)
 
           raise response.message unless response.success?
 
+          response.payload[:add_on_purchase].update!(users: [User.find_by_username('root')])
+
           puts "Code suggestions add-on added..."
         end
 

ee/lib/tasks/gitlab/duo.rake

@@ -3,7 +3,12 @@
 namespace :gitlab do
   namespace :duo do
     desc 'GitLab | Duo | Enable GitLab Duo features on the specified group'
-    task :setup, [:root_group_path] => :environment do |_, args|
+    task :setup, [:root_group_path, :add_on] => :environment do |_, args|
+      Gitlab::Duo::Developments::Setup.new(args).execute
+    end
+
+    desc 'GitLab | Duo | Enable GitLab Duo features for the instance'
+    task :setup_instance, [:add_on] => :environment do |_, args|
       Gitlab::Duo::Developments::Setup.new(args).execute
     end
 

ee/spec/lib/gitlab/duo/developments/setup_spec.rb

@@ -73,8 +73,11 @@
   end
 
   shared_examples 'creates add-on purchases' do
-    it 'creates add-on purchases', :aggregate_failures do
-      expect { setup }.to change { ::GitlabSubscriptions::AddOnPurchase.count }.by(2)
+    it 'creates enterprise add-on purchases', :aggregate_failures do
+      setup
+
+      expect(::GitlabSubscriptions::AddOnPurchase.for_gitlab_duo_pro.count).to eq(0)
+      expect(::GitlabSubscriptions::AddOnPurchase.for_duo_enterprise.count).to eq(1)
     end
   end
 
@@ -119,6 +122,17 @@
       end
     end
 
+    context 'when creating duo pro add on' do
+      let(:args) { { root_group_path: 'test', add_on: 'duo_pro' } }
+
+      it 'creates duo pro add-on only' do
+        setup
+
+        expect(::GitlabSubscriptions::AddOnPurchase.for_gitlab_duo_pro.count).to eq(1)
+        expect(::GitlabSubscriptions::AddOnPurchase.for_duo_enterprise.count).to eq(0)
+      end
+    end
+
     it_behaves_like 'checks for dev or test env'
     it_behaves_like 'errors when GITLAB_SIMULATE_SAAS has unexpected value', true
     it_behaves_like 'enables all necessary feature flags'