Moved user cap worker logic to User

As a result, this logic is now used  when activating a User upon SAML sign-in.
Tests related to these new methods in EE::User were also added.

Changelog: added
EE: true

ee/app/models/ee/user.rb

@@ -168,6 +168,16 @@ module EE
 
         scope
       end
+
+      def user_cap_reached?
+        return false unless user_cap_max
+
+        billable.limit(user_cap_max + 1).count >= user_cap_max
+      end
+
+      def user_cap_max
+        ::Gitlab::CurrentSettings.new_user_signups_cap
+      end
     end
 
     def cannot_be_admin_and_auditor
@@ -410,6 +420,16 @@ module EE
       has_valid_credit_card? || !requires_credit_card_to_enable_shared_runners?(project)
     end
 
+    def activate_based_on_user_cap?
+      !blocked_auto_created_omniauth_user? &&
+        blocked_pending_approval? &&
+        self.class.user_cap_max.present?
+    end
+
+    def blocked_auto_created_omniauth_user?
+      ::Gitlab.config.omniauth.block_auto_created_users && identities.any?
+    end
+
     protected
 
     override :password_required?
@@ -487,6 +507,7 @@ module EE
 
     def perform_user_cap_check
       return unless ::Gitlab::CurrentSettings.should_apply_user_signup_cap?
+      return if active?
 
       run_after_commit do
         SetUserStatusBasedOnUserCapSettingWorker.perform_async(id)

ee/app/workers/set_user_status_based_on_user_cap_setting_worker.rb

@@ -15,11 +15,9 @@ class SetUserStatusBasedOnUserCapSettingWorker
   def perform(user_id)
     user = User.includes(:identities).find_by_id(user_id) # rubocop: disable CodeReuse/ActiveRecord
 
-    return if blocked_auto_created_omniauth_user?(user)
-    return unless user.blocked_pending_approval?
-    return if user_cap_max.nil?
+    return unless user.activate_based_on_user_cap?
 
-    if user_cap_reached?
+    if User.user_cap_reached?
       send_user_cap_reached_email
       return
     end
@@ -37,27 +35,9 @@ class SetUserStatusBasedOnUserCapSettingWorker
 
   private
 
-  def user_cap_max
-    strong_memoize(:user_cap_max) do
-      ::Gitlab::CurrentSettings.new_user_signups_cap
-    end
-  end
-
-  def current_billable_users_count
-    User.billable.count
-  end
-
-  def user_cap_reached?
-    current_billable_users_count >= user_cap_max
-  end
-
   def send_user_cap_reached_email
     User.admins.active.each do |user|
       ::Notify.user_cap_reached(user.id).deliver_later
     end
   end
-
-  def blocked_auto_created_omniauth_user?(user)
-    ::Gitlab.config.omniauth.block_auto_created_users && user.identities.any?
-  end
 end

ee/lib/ee/gitlab/auth/saml/user.rb

@@ -13,6 +13,7 @@ module EE
 
             if user_in_required_group?
               unblock_user(user, "in required group") if user&.persisted? && user&.ldap_blocked?
+              unblock_user(user, "user cap not reached yet") if activate_user_based_on_user_cap?(user)
             elsif user&.persisted?
               block_user(user, "not in required group") unless user.blocked?
             else
@@ -31,6 +32,17 @@ module EE
 
           protected
 
+          def activate_user_based_on_user_cap?(user)
+            return false unless user&.activate_based_on_user_cap?
+
+            begin
+              !::User.user_cap_reached?
+            rescue ActiveRecord::QueryAborted => e
+              ::Gitlab::ErrorTracking.track_exception(e, saml_user_email: user.email)
+              false
+            end
+          end
+
           def block_user(user, reason)
             user.ldap_block
             log_user_changes(user, "#{reason}, blocking")

ee/spec/lib/gitlab/auth/saml/user_spec.rb

@@ -144,6 +144,67 @@ RSpec.describe Gitlab::Auth::Saml::User do
             expect(saml_user.find_user).to be_ldap_blocked
           end
         end
+
+        context 'when a sign-up user cap has been set' do
+          before do
+            saml_user.gl_user.state = ::User::BLOCKED_PENDING_APPROVAL_STATE
+            stub_application_setting(new_user_signups_cap: new_user_signups_cap)
+          end
+
+          context 'when the user cap has been reached' do
+            let(:new_user_signups_cap) { 1 }
+
+            it 'does not activate the user' do
+              create(:user)
+
+              saml_user.save
+
+              expect(saml_user.find_user).to be_blocked
+            end
+          end
+
+          context 'when the user cap has not been reached' do
+            let(:new_user_signups_cap) { 100 }
+
+            before do
+              stub_omniauth_setting(block_auto_created_users: block)
+            end
+
+            context 'when the user can be activated based on user cap' do
+              let(:block) { false }
+
+              it 'activates the user' do
+                saml_user.save
+
+                expect(saml_user.find_user).to be_active
+              end
+
+              context 'when the query behind .user_cap_reached? times out' do
+                it 'does not activate the user' do
+                  allow(::User).to receive(:user_cap_reached?).and_raise(ActiveRecord::QueryAborted)
+
+                  saml_user.save
+
+                  expect(::Gitlab::ErrorTracking).to receive(:track_exception).with(
+                    instance_of(ActiveRecord::QueryAborted),
+                    saml_user_email: saml_user.gl_user.email
+                  )
+                  expect(saml_user.find_user).to be_blocked
+                end
+              end
+            end
+
+            context 'when the user cannot be activated based on user cap' do
+              let(:block) { true }
+
+              it 'does not activate the user' do
+                saml_user.save
+
+                expect(saml_user.find_user).to be_blocked
+              end
+            end
+          end
+        end
       end
     end
   end

ee/spec/models/ee/user_spec.rb

@@ -129,6 +129,14 @@ RSpec.describe User do
 
           create(:user, state: 'blocked_pending_approval')
         end
+
+        context 'when the user is already active' do
+          it 'does not enqueue SetUserStatusBasedOnUserCapSettingWorker' do
+            expect(SetUserStatusBasedOnUserCapSettingWorker).not_to receive(:perform_async)
+
+            create(:user, state: 'active')
+          end
+        end
       end
     end
   end
@@ -1936,4 +1944,88 @@ RSpec.describe User do
       expect(user.escalation_policies).to contain_exactly(policy)
     end
   end
+
+  describe '.user_cap_reached?' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { described_class.user_cap_reached? }
+
+    where(:billable_count, :user_cap_max, :result) do
+      2 | nil | false
+      2 | 5   | false
+      5 | 5   | true
+      8 | 5   | true
+    end
+
+    with_them do
+      before do
+        allow(described_class).to receive_message_chain(:billable, :limit).and_return(Array.new(billable_count, instance_double('User')))
+        allow(described_class).to receive(:user_cap_max).and_return(user_cap_max)
+      end
+
+      it { is_expected.to eq(result) }
+    end
+  end
+
+  describe '.user_cap_max' do
+    it 'is equal to new_user_signups_cap setting' do
+      cap = 10
+      stub_application_setting(new_user_signups_cap: cap)
+
+      expect(described_class.user_cap_max).to eq(cap)
+    end
+  end
+
+  describe '#blocked_auto_created_omniauth_user?' do
+    context 'when the auto-creation of an omniauth user is blocked' do
+      before do
+        stub_omniauth_setting(block_auto_created_users: true)
+      end
+
+      context 'when the user is an omniauth user' do
+        it 'is true' do
+          omniauth_user = create(:omniauth_user)
+
+          expect(omniauth_user.blocked_auto_created_omniauth_user?).to be_truthy
+        end
+      end
+
+      context 'when the user is not an omniauth user' do
+        it 'is false' do
+          user = build(:user)
+
+          expect(user.blocked_auto_created_omniauth_user?).to be_falsey
+        end
+      end
+    end
+  end
+
+  describe '#activate_based_on_user_cap?' do
+    using RSpec::Parameterized::TableSyntax
+
+    let_it_be(:user) { create(:user) }
+
+    subject { user.activate_based_on_user_cap? }
+
+    where(:blocked_auto_created_omniauth, :blocked_pending_approval, :user_cap_max_present, :result) do
+      true  | true  | true  | false
+      false | true  | true  | true
+      true  | false | true  | false
+      false | false | true  | false
+      true  | true  | false | false
+      false | true  | false | false
+      true  | false | false | false
+      false | false | false | false
+    end
+
+    with_them do
+      before do
+        allow(user).to receive(:blocked_auto_created_omniauth_user?).and_return(blocked_auto_created_omniauth)
+        allow(user).to receive(:blocked_pending_approval?).and_return(blocked_pending_approval)
+        allow(described_class.user_cap_max).to receive(:present?).and_return(user_cap_max_present)
+      end
+
+      it { is_expected.to eq(result) }
+    end
+  end
 end