Extend bulk_user_assignment lib

Extend the lib to support gitlab.com

ee/lib/duo_pro/bulk_user_assignment.rb

@@ -17,6 +17,7 @@
 
 module DuoPro
   class BulkUserAssignment
+    include ::GitlabSubscriptions::SubscriptionHelper
     attr_reader :usernames, :add_on_purchase, :successful_assignments, :failed_assignments
 
     def initialize(usernames, add_on_purchase)
@@ -57,10 +58,13 @@ def process_users(usernames)
     end
 
     def assign(user)
-      ::GitlabSubscriptions::UserAddOnAssignments::SelfManaged::CreateService.new(
-        add_on_purchase: add_on_purchase,
-        user: user
-      ).execute
+      service_class = if gitlab_com_subscription?
+                        ::GitlabSubscriptions::UserAddOnAssignments::Saas::CreateService
+                      else
+                        ::GitlabSubscriptions::UserAddOnAssignments::SelfManaged::CreateService
+                      end
+
+      service_class.new(add_on_purchase: add_on_purchase, user: user).execute
     end
 
     def log_no_seats_available(result, username)

ee/spec/lib/duo_pro/bulk_user_assignment_spec.rb

@@ -15,19 +15,46 @@
   end
 
   describe '#execute' do
-    let(:add_on) { create(:gitlab_subscription_add_on) }
+    let_it_be(:add_on) { create(:gitlab_subscription_add_on) }
     let(:usernames) { User.pluck(:username) + ['code_suggestions_not_found_username'] }
 
-    before do
-      create(:user, username: 'code_suggestions_active_user1')
-      create(:user, username: 'code_suggestions_active_user2')
-      create(:user, username: 'code_suggestions_active_user3')
-      create(:user, username: 'code_suggestions_extra_user1')
-      create(:user, username: 'code_suggestions_extra_user2')
-      create(:user, :blocked, username: 'code_suggestions_blocked_user')
-      create(:user, :banned, username: 'code_suggestions_banned_user')
-      create(:user, :bot, username: 'code_suggestions_bot_user')
-      create(:user, :ghost, username: 'code_suggestions_ghost_user')
+    subject(:bulk_assignment) { described_class.new(usernames, add_on_purchase) }
+
+    shared_examples 'bulk user assignment with enough seats' do
+      it 'returns success and failed assignments' do
+        results = bulk_assignment.execute
+
+        expect(results[:successful_assignments]).to eq([
+          "User assigned: code_suggestions_active_user1",
+          "User assigned: code_suggestions_active_user2",
+          "User assigned: code_suggestions_active_user3",
+          "User assigned: code_suggestions_extra_user1",
+          "User assigned: code_suggestions_extra_user2"
+        ])
+
+        expect(results[:failed_assignments]).to eq([
+          "Failed to assign seat to user: code_suggestions_blocked_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
+          "Failed to assign seat to user: code_suggestions_banned_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
+          "Failed to assign seat to user: code_suggestions_bot_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
+          "Failed to assign seat to user: code_suggestions_ghost_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
+          "User is not found: code_suggestions_not_found_username"
+        ])
+      end
+    end
+
+    shared_examples 'bulk user assignment with not enough seats' do
+      it 'returns success and failed assignments and stops execution' do
+        results = bulk_assignment.execute
+
+        expect(results[:successful_assignments]).to eq(
+          ["User assigned: code_suggestions_active_user1",
+            "User assigned: code_suggestions_active_user2",
+            "User assigned: code_suggestions_active_user3"])
+
+        expect(results[:failed_assignments]).to eq(
+          ["Failed to assign seat to user: code_suggestions_extra_user1, Errors: [\"NO_SEATS_AVAILABLE\"]",
+            "##No seats are left; users starting from @code_suggestions_extra_user1 onwards were not assigned.##"])
+      end
     end
 
     context 'when the AddOn is not purchased' do
@@ -37,51 +64,72 @@
       end
     end
 
-    context 'when the AddOn is purchased' do
-      let(:add_on_purchase) do
-        create(:gitlab_subscription_add_on_purchase, :self_managed, quantity: 10, add_on: add_on)
+    context 'on self managed instances' do
+      before do
+        create(:user, username: 'code_suggestions_active_user1')
+        create(:user, username: 'code_suggestions_active_user2')
+        create(:user, username: 'code_suggestions_active_user3')
+        create(:user, username: 'code_suggestions_extra_user1')
+        create(:user, username: 'code_suggestions_extra_user2')
+        create(:user, :blocked, username: 'code_suggestions_blocked_user')
+        create(:user, :banned, username: 'code_suggestions_banned_user')
+        create(:user, :bot, username: 'code_suggestions_bot_user')
+        create(:user, :ghost, username: 'code_suggestions_ghost_user')
       end
 
-      subject(:bulk_assignment) { described_class.new(usernames, add_on_purchase) }
-
-      context 'with enough seats' do
-        it 'returns success and failed assignments' do
-          results = bulk_assignment.execute
+      context 'when the AddOn is purchased' do
+        context 'with enough seats' do
+          include_examples 'bulk user assignment with enough seats' do
+            let(:add_on_purchase) do
+              create(:gitlab_subscription_add_on_purchase, :self_managed, quantity: 10, add_on: add_on)
+            end
+          end
+        end
 
-          expect(results[:successful_assignments]).to eq([
-            "User assigned: code_suggestions_active_user1",
-            "User assigned: code_suggestions_active_user2",
-            "User assigned: code_suggestions_active_user3",
-            "User assigned: code_suggestions_extra_user1",
-            "User assigned: code_suggestions_extra_user2"
-          ])
-
-          expect(results[:failed_assignments]).to eq([
-            "Failed to assign seat to user: code_suggestions_blocked_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
-            "Failed to assign seat to user: code_suggestions_banned_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
-            "Failed to assign seat to user: code_suggestions_bot_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
-            "Failed to assign seat to user: code_suggestions_ghost_user, Errors: [\"INVALID_USER_MEMBERSHIP\"]",
-            "User is not found: code_suggestions_not_found_username"
-          ])
+        context 'with not enough seats' do
+          include_examples 'bulk user assignment with not enough seats' do
+            let(:add_on_purchase) do
+              create(:gitlab_subscription_add_on_purchase, :self_managed, quantity: 3, add_on: add_on)
+            end
+          end
         end
       end
+    end
 
-      context 'with not enough seats' do
-        let(:add_on_purchase) do
-          create(:gitlab_subscription_add_on_purchase, :self_managed, quantity: 3, add_on: add_on)
+    context 'on Gitlab.com' do
+      before do
+        stub_saas_features(gitlab_com_subscriptions: true)
+      end
+
+      let_it_be(:namespace) { create(:group) }
+
+      context 'with bulk assignment' do
+        before_all do
+          namespace.add_developer(create(:user, username: 'code_suggestions_active_user1'))
+          namespace.add_developer(create(:user, username: 'code_suggestions_active_user2'))
+          namespace.add_developer(create(:user, username: 'code_suggestions_active_user3'))
+          namespace.add_developer(create(:user, username: 'code_suggestions_extra_user1'))
+          namespace.add_developer(create(:user, username: 'code_suggestions_extra_user2'))
+          namespace.add_developer(create(:user, :blocked, username: 'code_suggestions_blocked_user'))
+          namespace.add_developer(create(:user, :banned, username: 'code_suggestions_banned_user'))
+          namespace.add_developer(create(:user, :bot, username: 'code_suggestions_bot_user'))
+          namespace.add_developer(create(:user, :ghost, username: 'code_suggestions_ghost_user'))
         end
 
-        it 'returns success and failed assignments and stops execution' do
-          results = bulk_assignment.execute
+        context 'with enough seats' do
+          let_it_be(:add_on_purchase) do
+            create(:gitlab_subscription_add_on_purchase, quantity: 10, namespace: namespace, add_on: add_on)
+          end
+
+          include_examples 'bulk user assignment with enough seats'
+        end
 
-          expect(results[:successful_assignments]).to eq(
-            ["User assigned: code_suggestions_active_user1",
-              "User assigned: code_suggestions_active_user2",
-              "User assigned: code_suggestions_active_user3"])
+        context 'with not enough seats' do
+          let_it_be(:add_on_purchase) do
+            create(:gitlab_subscription_add_on_purchase, quantity: 3, namespace: namespace, add_on: add_on)
+          end
 
-          expect(results[:failed_assignments]).to eq(
-            ["Failed to assign seat to user: code_suggestions_extra_user1, Errors: [\"NO_SEATS_AVAILABLE\"]",
-              "##No seats are left; users starting from @code_suggestions_extra_user1 onwards were not assigned.##"])
+          include_examples 'bulk user assignment with not enough seats'
         end
       end
     end