Use excess storage size logic for storage banners

Implement a toggle between using the first version of storage banners
and the new version.

ee/app/graphql/ee/types/namespace_type.rb

@@ -43,7 +43,7 @@ module EE
               GraphQL::FLOAT_TYPE,
               null: true,
               description: 'Total storage limit of the root namespace in bytes',
-              resolve: -> (obj, _args, _ctx) { EE::Namespace::RootStorageSize.new(obj).limit }
+              resolve: -> (obj, _args, _ctx) { obj.root_storage_size.limit }
 
         field :is_temporary_storage_increase_enabled,
               GraphQL::BOOLEAN_TYPE,

ee/app/helpers/ee/namespace_storage_limit_alert_helper.rb

@@ -16,7 +16,7 @@ module EE
     def namespace_storage_alert(namespace)
       return {} if current_user.nil?
 
-      payload = Namespaces::CheckStorageSizeService.new(namespace, current_user).execute.payload
+      payload = check_storage_size_service(namespace).execute.payload
 
       return {} if payload.empty?
 
@@ -70,5 +70,15 @@ module EE
 
       EE::SUBSCRIPTIONS_MORE_STORAGE_URL
     end
+
+    private
+
+    def check_storage_size_service(namespace)
+      if namespace.additional_repo_storage_by_namespace_enabled?
+        ::Namespaces::CheckExcessStorageSizeService.new(namespace, current_user)
+      else
+        ::Namespaces::CheckStorageSizeService.new(namespace, current_user)
+      end
+    end
   end
 end

ee/app/models/ee/namespace.rb

@@ -193,8 +193,8 @@ module EE
 
     def over_storage_limit?
       ::Gitlab::CurrentSettings.enforce_namespace_storage_limit? &&
-        ::Feature.enabled?(:namespace_storage_limit, root_ancestor) &&
-        RootStorageSize.new(root_ancestor).above_size_limit?
+      ::Feature.enabled?(:namespace_storage_limit, root_ancestor) &&
+        root_ancestor.root_storage_size.above_size_limit?
     end
 
     def total_repository_size_excess
@@ -386,6 +386,15 @@ module EE
       update(temporary_storage_increase_ends_on: TEMPORARY_STORAGE_INCREASE_DAYS.days.from_now)
     end
 
+    def additional_repo_storage_by_namespace_enabled?
+      !::Feature.enabled?(:namespace_storage_limit, self) && ::Feature.enabled?(:additional_repo_storage_by_namespace, self)
+    end
+
+    def root_storage_size
+      klass = additional_repo_storage_by_namespace_enabled? ? RootExcessStorageSize : RootStorageSize
+      klass.new(self)
+    end
+
     private
 
     def fallback_plan

ee/app/models/namespace_limit.rb

@@ -23,7 +23,7 @@ class NamespaceLimit < ApplicationRecord
   def eligible_for_temporary_storage_increase?
     return false unless ::Feature.enabled?(:temporary_storage_increase, namespace)
 
-    EE::Namespace::RootStorageSize.new(namespace).usage_ratio >= MIN_REQURIED_STORAGE_USAGE_RATIO
+    namespace.root_storage_size.usage_ratio >= MIN_REQURIED_STORAGE_USAGE_RATIO
   end
 
   private

ee/app/services/ee/post_receive_service.rb

@@ -56,12 +56,20 @@ module EE
     def storage_size_limit_alert
       return unless repository&.repo_type&.project?
 
-      payload = Namespaces::CheckStorageSizeService.new(project.namespace, user).execute.payload
+      payload = check_storage_size_service(project.namespace).execute.payload
       return unless payload.present?
 
       alert_level = "##### #{payload[:alert_level].to_s.upcase} #####"
 
       [alert_level, payload[:usage_message], payload[:explanation_message]].join("\n")
     end
+
+    def check_storage_size_service(namespace)
+      if namespace.additional_repo_storage_by_namespace_enabled?
+        ::Namespaces::CheckExcessStorageSizeService.new(namespace, user)
+      else
+        ::Namespaces::CheckStorageSizeService.new(namespace, user)
+      end
+    end
   end
 end

ee/spec/helpers/ee/namespace_storage_limit_alert_helper_spec.rb

@@ -71,47 +71,59 @@ RSpec.describe EE::NamespaceStorageLimitAlertHelper do
       }
     end
 
-    before do
-      allow(helper).to receive(:current_user).and_return(admin)
-      allow_next_instance_of(Namespaces::CheckStorageSizeService, namespace, admin) do |check_storage_size_service|
-        expect(check_storage_size_service).to receive(:execute).and_return(ServiceResponse.success(payload: payload))
-      end
-    end
-
-    context 'when payload is not empty and no cookie is set' do
-      it { is_expected.to eq(payload) }
+    where(:namespace_storage_limit_enabled, :additional_repo_storage_by_namespace_enabled, :service_class_name) do
+      true  | false | Namespaces::CheckStorageSizeService
+      true  | true  | Namespaces::CheckStorageSizeService
+      false | true  | Namespaces::CheckExcessStorageSizeService
+      false | false | Namespaces::CheckStorageSizeService
     end
 
-    context 'when there is no current_user' do
+    with_them do
       before do
-        allow(helper).to receive(:current_user).and_return(nil)
-      end
+        stub_feature_flags(namespace_storage_limit: namespace_storage_limit_enabled)
+        stub_feature_flags(additional_repo_storage_by_namespace: additional_repo_storage_by_namespace_enabled)
 
-      it { is_expected.to eq({}) }
-    end
+        allow(helper).to receive(:current_user).and_return(admin)
+        allow_next_instance_of(service_class_name, namespace, admin) do |service|
+          expect(service).to receive(:execute).and_return(ServiceResponse.success(payload: payload))
+        end
+      end
 
-    context 'when payload is empty' do
-      let(:payload) { {} }
+      context 'when payload is not empty and no cookie is set' do
+        it { is_expected.to eq(payload) }
+      end
 
-      it { is_expected.to eq({}) }
-    end
+      context 'when there is no current_user' do
+        before do
+          allow(helper).to receive(:current_user).and_return(nil)
+        end
 
-    context 'when cookie is set' do
-      before do
-        helper.request.cookies["hide_storage_limit_alert_#{namespace.id}_info"] = 'true'
+        it { is_expected.to eq({}) }
       end
 
-      it { is_expected.to eq({}) }
-    end
+      context 'when payload is empty' do
+        let(:payload) { {} }
 
-    context 'when payload is empty and cookie is set' do
-      let(:payload) { {} }
+        it { is_expected.to eq({}) }
+      end
 
-      before do
-        helper.request.cookies["hide_storage_limit_alert_#{namespace.id}_info"] = 'true'
+      context 'when cookie is set' do
+        before do
+          helper.request.cookies["hide_storage_limit_alert_#{namespace.id}_info"] = 'true'
+        end
+
+        it { is_expected.to eq({}) }
       end
 
-      it { is_expected.to eq({}) }
+      context 'when payload is empty and cookie is set' do
+        let(:payload) { {} }
+
+        before do
+          helper.request.cookies["hide_storage_limit_alert_#{namespace.id}_info"] = 'true'
+        end
+
+        it { is_expected.to eq({}) }
+      end
     end
   end
 

ee/spec/models/ee/namespace_limit_spec.rb

@@ -5,12 +5,17 @@ require 'spec_helper'
 RSpec.describe NamespaceLimit do
   let(:namespace_limit) { build(:namespace_limit) }
   let(:usage_ratio) { 0.5 }
+  let(:namespace_storage_limit_enabled) { true }
 
   subject { namespace_limit }
 
   before do
-    allow_next_instance_of(EE::Namespace::RootStorageSize, namespace_limit.namespace) do |root_storage|
-      allow(root_storage).to receive(:usage_ratio).and_return(usage_ratio)
+    stub_feature_flags(namespace_storage_limit: namespace_storage_limit_enabled)
+
+    [EE::Namespace::RootStorageSize, EE::Namespace::RootExcessStorageSize].each do |class_name|
+      allow_next_instance_of(class_name, namespace_limit.namespace) do |root_storage|
+        allow(root_storage).to receive(:usage_ratio).and_return(usage_ratio)
+      end
     end
   end
 
@@ -54,21 +59,39 @@ RSpec.describe NamespaceLimit do
     context 'when usage ratio is above the threshold' do
       let(:usage_ratio) { 0.5 }
 
-      it { is_expected.to be_truthy }
+      it { is_expected.to eq(true) }
 
-      context 'when feature is disabled' do
+      context 'when feature flag :temporary_storage_increase disabled' do
         before do
           stub_feature_flags(temporary_storage_increase: false)
         end
 
+        context 'when feature flag :namespace_storage_limit disabled' do
+          let(:namespace_storage_limit_enabled) { false }
+
+          it { is_expected.to eq(false) }
+        end
+
         it { is_expected.to eq(false) }
       end
+
+      context 'when feature flag :namespace_storage_limit disabled' do
+        let(:namespace_storage_limit_enabled) { false }
+
+        it { is_expected.to eq(true) }
+      end
     end
 
     context 'when usage ratio is below the threshold' do
       let(:usage_ratio) { 0.49 }
 
-      it { is_expected.to be_falsey }
+      context 'when feature flag :namespace_storage_limit disabled' do
+        let(:namespace_storage_limit_enabled) { false }
+
+        it { is_expected.to eq(false) }
+      end
+
+      it { is_expected.to eq(false) }
     end
   end
 
@@ -126,6 +149,12 @@ RSpec.describe NamespaceLimit do
         let(:usage_ratio) { 0.5 }
 
         it { is_expected.to be_valid }
+
+        context 'when feature flag :namespace_storage_limit disabled' do
+          let(:namespace_storage_limit_enabled) { false }
+
+          it { is_expected.to be_valid }
+        end
       end
 
       context 'when storage usage is below threshold' do
@@ -135,6 +164,15 @@ RSpec.describe NamespaceLimit do
           expect(namespace_limit).to be_invalid
           expect(namespace_limit.errors[:temporary_storage_increase_ends_on]).to include("can only be set with more than 50% usage")
         end
+
+        context 'when feature flag :namespace_storage_limit disabled' do
+          let(:namespace_storage_limit_enabled) { false }
+
+          it 'is invalid' do
+            expect(namespace_limit).to be_invalid
+            expect(namespace_limit.errors[:temporary_storage_increase_ends_on]).to include("can only be set with more than 50% usage")
+          end
+        end
       end
     end
   end

ee/spec/models/ee/namespace_spec.rb

@@ -1770,6 +1770,32 @@ RSpec.describe Namespace do
     end
   end
 
+  describe '#root_storage_size' do
+    let_it_be(:namespace) { build(:namespace) }
+
+    subject { namespace.root_storage_size }
+
+    context 'with feature flag :namespace_storage_limit enabled' do
+      it 'initializes a new instance of EE::Namespace::RootStorageSize' do
+        expect(EE::Namespace::RootStorageSize).to receive(:new).with(namespace)
+
+        subject
+      end
+    end
+
+    context 'with feature flag :namespace_storage_limit disabled' do
+      before do
+        stub_feature_flags(namespace_storage_limit: false)
+      end
+
+      it 'initializes a new instance of EE::Namespace::RootExcessStorageSize' do
+        expect(EE::Namespace::RootExcessStorageSize).to receive(:new).with(namespace)
+
+        subject
+      end
+    end
+  end
+
   def create_project(repository_size:, lfs_objects_size:, repository_size_limit:)
     create(:project, namespace: namespace, repository_size_limit: repository_size_limit).tap do |project|
       create(:project_statistics, project: project, repository_size: repository_size, lfs_objects_size: lfs_objects_size)

ee/spec/services/ee/post_receive_service_spec.rb

@@ -104,36 +104,50 @@ RSpec.describe PostReceiveService, :geo do
   end
 
   describe 'storage size limit alerts' do
+    using RSpec::Parameterized::TableSyntax
+
     let(:check_storage_size_response) { ServiceResponse.success }
 
-    before do
-      expect_next_instance_of(Namespaces::CheckStorageSizeService, project.namespace, user) do |check_storage_size_service|
-        expect(check_storage_size_service).to receive(:execute).and_return(check_storage_size_response)
-      end
+    where(:namespace_storage_limit_enabled, :additional_repo_storage_by_namespace_enabled, :service_class_name) do
+      true  | false | Namespaces::CheckStorageSizeService
+      true  | true  | Namespaces::CheckStorageSizeService
+      false | true  | Namespaces::CheckExcessStorageSizeService
+      false | false | Namespaces::CheckStorageSizeService
     end
 
-    context 'when there is no payload' do
-      it 'adds no alert' do
-        expect(subject.size).to eq(0)
+    with_them do
+      before do
+        stub_feature_flags(namespace_storage_limit: namespace_storage_limit_enabled)
+        stub_feature_flags(additional_repo_storage_by_namespace: additional_repo_storage_by_namespace_enabled)
+
+        allow_next_instance_of(service_class_name, project.namespace, user) do |service|
+          expect(service).to receive(:execute).and_return(check_storage_size_response)
+        end
       end
-    end
 
-    context 'when there is payload' do
-      let(:check_storage_size_response) do
-        ServiceResponse.success(
-          payload: {
-            alert_level: :info,
-            usage_message: "Usage",
-            explanation_message: "Explanation"
-          }
-        )
+      context 'when there is no payload' do
+        it 'adds no alert' do
+          expect(subject).to be_empty
+        end
       end
 
-      it 'adds an alert' do
-        response = subject
+      context 'when there is payload' do
+        let(:check_storage_size_response) do
+          ServiceResponse.success(
+            payload: {
+              alert_level: :info,
+              usage_message: "Usage",
+              explanation_message: "Explanation"
+            }
+          )
+        end
 
-        expect(response.size).to eq(1)
-        expect(response).to include({ 'type' => 'alert', 'message' => "##### INFO #####\nUsage\nExplanation" })
+        it 'adds an alert' do
+          response = subject
+
+          expect(response).to be_present
+          expect(response).to include({ 'type' => 'alert', 'message' => "##### INFO #####\nUsage\nExplanation" })
+        end
       end
     end
   end