Check if feature flag is enabled for namespace

The feature flag :additional_repo_storage_by_namespace check in
EE::Gitlab::RepositorySizeChecker was done globally. Now it will check
if it's enabled for the given namespace.

app/models/snippet.rb

@@ -283,8 +283,7 @@ class Snippet < ApplicationRecord
       ::Gitlab::RepositorySizeChecker.new(
         current_size_proc: -> { repository.size.megabytes },
         limit: Gitlab::CurrentSettings.snippet_size_limit,
-        total_repository_size_excess: nil,
-        additional_purchased_storage: nil
+        namespace: nil
       )
     end
   end

ee/app/models/ee/project.rb

@@ -532,8 +532,7 @@ module EE
         ::Gitlab::RepositorySizeChecker.new(
           current_size_proc: -> { statistics.total_repository_size },
           limit: actual_size_limit,
-          total_repository_size_excess: namespace.total_repository_size_excess,
-          additional_purchased_storage: namespace.additional_purchased_storage_size.megabytes,
+          namespace: namespace,
           enabled: License.feature_available?(:repository_size_limit)
         )
       end

ee/app/models/ee/snippet.rb

@@ -16,8 +16,7 @@ module EE
         ::Gitlab::RepositorySizeChecker.new(
           current_size_proc: -> { repository.size.megabytes },
           limit: ::Gitlab::CurrentSettings.snippet_size_limit,
-          total_repository_size_excess: project&.namespace&.total_repository_size_excess,
-          additional_purchased_storage: project&.namespace&.additional_purchased_storage_size&.megabytes
+          namespace: project&.namespace
         )
       end
     end

ee/lib/ee/gitlab/repository_size_checker.rb

@@ -27,7 +27,15 @@ module EE
       def additional_repo_storage_available?
         return false unless ::Gitlab::CurrentSettings.automatic_purchased_storage_allocation?
 
-        ::Feature.enabled?(:additional_repo_storage_by_namespace)
+        ::Feature.enabled?(:additional_repo_storage_by_namespace, namespace)
+      end
+
+      def total_repository_size_excess
+        namespace&.total_repository_size_excess.to_i
+      end
+
+      def additional_purchased_storage
+        namespace&.additional_purchased_storage_size&.megabytes.to_i
       end
 
       def remaining_additional_purchased_storage

ee/spec/lib/ee/gitlab/repository_size_checker_spec.rb

@@ -5,6 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::RepositorySizeChecker do
   let(:current_size) { 0 }
   let(:limit) { 50 }
+  let(:namespace) { build(:namespace, additional_purchased_storage_size: additional_purchased_storage) }
   let(:total_repository_size_excess) { 0 }
   let(:additional_purchased_storage) { 0 }
   let(:enabled) { true }
@@ -12,16 +13,17 @@ RSpec.describe Gitlab::RepositorySizeChecker do
 
   subject do
     described_class.new(
-      current_size_proc: -> { current_size },
-      limit: limit,
-      total_repository_size_excess: total_repository_size_excess,
-      additional_purchased_storage: additional_purchased_storage,
+      current_size_proc: -> { current_size.megabytes },
+      limit: limit.megabytes,
+      namespace: namespace,
       enabled: enabled
     )
   end
 
   before do
     allow(Gitlab::CurrentSettings).to receive(:automatic_purchased_storage_allocation?).and_return(gitlab_setting_enabled)
+
+    allow(namespace).to receive(:total_repository_size_excess).and_return(total_repository_size_excess.megabytes)
   end
 
   describe '#above_size_limit?' do
@@ -33,24 +35,55 @@ RSpec.describe Gitlab::RepositorySizeChecker do
     end
 
     context 'with feature flag :additional_repo_storage_by_namespace enabled' do
-      context 'when there is available excess storage' do
-        it 'returns false' do
+      shared_examples 'feature flag :additional_repo_storage_by_namespace enabled' do
+        context 'when there is available excess storage' do
+          it 'returns false' do
+            expect(subject.above_size_limit?).to eq(false)
+          end
+        end
+
+        context 'when size is above the limit and there is no excess storage' do
+          let(:current_size) { 100 }
+          let(:total_repository_size_excess) { 20 }
+          let(:additional_purchased_storage) { 10 }
+
+          it 'returns true' do
+            expect(subject.above_size_limit?).to eq(true)
+          end
+        end
+
+        it 'returns false when not over the limit' do
           expect(subject.above_size_limit?).to eq(false)
         end
       end
 
-      context 'when size is above the limit and there is no exccess storage' do
-        let(:current_size) { 100 }
-        let(:total_repository_size_excess) { 20 }
-        let(:additional_purchased_storage) { 10 }
+      include_examples 'feature flag :additional_repo_storage_by_namespace enabled'
 
-        it 'returns true' do
-          expect(subject.above_size_limit?).to eq(true)
+      context 'when only enabled for a specific namespace' do
+        before do
+          stub_feature_flags(additional_repo_storage_by_namespace: false)
         end
-      end
 
-      it 'returns false when not over the limit' do
-        expect(subject.above_size_limit?).to eq(false)
+        context 'when namespace is another one than the given one' do
+          let(:other_namespace) { create(:namespace) }
+
+          before do
+            stub_feature_flags(additional_repo_storage_by_namespace: other_namespace)
+          end
+
+          include_examples 'checker size above limit'
+          include_examples 'checker size not over limit'
+        end
+
+        context 'when namespace is the given one' do
+          let(:namespace) { create(:namespace, additional_purchased_storage_size: additional_purchased_storage) }
+
+          before do
+            stub_feature_flags(additional_repo_storage_by_namespace: namespace)
+          end
+
+          include_examples 'feature flag :additional_repo_storage_by_namespace enabled'
+        end
       end
     end
 
@@ -88,7 +121,7 @@ RSpec.describe Gitlab::RepositorySizeChecker do
         let(:additional_purchased_storage) { 10 }
 
         it 'returns 1' do
-          expect(subject.exceeded_size).to eq(1)
+          expect(subject.exceeded_size).to eq(1.megabytes)
         end
       end
 
@@ -96,7 +129,7 @@ RSpec.describe Gitlab::RepositorySizeChecker do
         let(:current_size) { 50 }
 
         it 'returns 1' do
-          expect(subject.exceeded_size(1)).to eq(1)
+          expect(subject.exceeded_size(1.megabytes)).to eq(1.megabytes)
         end
       end
 
@@ -104,7 +137,7 @@ RSpec.describe Gitlab::RepositorySizeChecker do
         let(:current_size) { 49 }
 
         it 'returns zero' do
-          expect(subject.exceeded_size(1)).to eq(0)
+          expect(subject.exceeded_size(1.megabytes)).to eq(0)
         end
       end
     end

ee/spec/models/project_spec.rb

@@ -2245,22 +2245,6 @@ RSpec.describe Project do
       end
     end
 
-    describe '#total_repository_size_excess' do
-      it 'returns the total repository size excess of the namespace' do
-        allow(project.namespace).to receive(:total_repository_size_excess).and_return(50)
-
-        expect(checker.total_repository_size_excess).to eq(50)
-      end
-    end
-
-    describe '#additional_purchased_storage' do
-      it 'returns the additional purchased storage size of the namespace' do
-        allow(project.namespace).to receive(:additional_purchased_storage_size).and_return(100)
-
-        expect(checker.additional_purchased_storage).to eq(100.megabytes)
-      end
-    end
-
     describe '#enabled?' do
       it 'returns true when not equal to zero' do
         project.repository_size_limit = 1

ee/spec/models/snippet_spec.rb

@@ -14,21 +14,14 @@ RSpec.describe Snippet do
     context 'when snippet belongs to a project' do
       subject { build(:project_snippet, project: project) }
 
-      let(:namespace) { build(:namespace, additional_purchased_storage_size: 50) }
+      let(:namespace) { build(:namespace) }
       let(:project) { build(:project, namespace: namespace) }
-      let(:total_repository_size_excess) { 100 }
-      let(:additional_purchased_storage) { 50.megabytes }
-
-      before do
-        allow(namespace).to receive(:total_repository_size_excess).and_return(total_repository_size_excess)
-      end
 
       include_examples 'size checker for snippet'
     end
 
     context 'when snippet without a project' do
-      let(:total_repository_size_excess) { 0 }
-      let(:additional_purchased_storage) { 0 }
+      let(:namespace) { nil }
 
       include_examples 'size checker for snippet'
     end

lib/gitlab/repository_size_checker.rb

@@ -3,14 +3,13 @@
 module Gitlab
   # Centralized class for repository size related calculations.
   class RepositorySizeChecker
-    attr_reader :limit, :total_repository_size_excess, :additional_purchased_storage
+    attr_reader :limit
 
     # @param current_size_proc [Proc] returns repository size in bytes
-    def initialize(current_size_proc:, limit:, total_repository_size_excess:, additional_purchased_storage:, enabled: true)
+    def initialize(current_size_proc:, limit:, namespace:, enabled: true)
       @current_size_proc = current_size_proc
       @limit = limit
-      @total_repository_size_excess = total_repository_size_excess.to_i
-      @additional_purchased_storage = additional_purchased_storage.to_i
+      @namespace = namespace
       @enabled = enabled && limit != 0
     end
 
@@ -44,6 +43,10 @@ module Gitlab
     def error_message
       @error_message_object ||= ::Gitlab::RepositorySizeErrorMessage.new(self)
     end
+
+    private
+
+    attr_reader :namespace
   end
 end
 

lib/gitlab/repository_size_error_message.rb

@@ -4,7 +4,7 @@ module Gitlab
   class RepositorySizeErrorMessage
     include ActiveSupport::NumberHelper
 
-    delegate :current_size, :limit, :total_repository_size_excess, :additional_purchased_storage, :exceeded_size, to: :@checker
+    delegate :current_size, :limit, :exceeded_size, to: :@checker
 
     # @param checher [RepositorySizeChecker]
     def initialize(checker)

spec/lib/gitlab/repository_size_checker_spec.rb

@@ -3,16 +3,16 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::RepositorySizeChecker do
+  let_it_be(:namespace) { nil }
   let(:current_size) { 0 }
   let(:limit) { 50 }
   let(:enabled) { true }
 
   subject do
     described_class.new(
-      current_size_proc: -> { current_size },
-      limit: limit,
-      total_repository_size_excess: 0,
-      additional_purchased_storage: 0,
+      current_size_proc: -> { current_size.megabytes },
+      limit: limit.megabytes,
+      namespace: namespace,
       enabled: enabled
     )
   end
@@ -20,7 +20,7 @@ RSpec.describe Gitlab::RepositorySizeChecker do
   describe '#enabled?' do
     context 'when enabled' do
       it 'returns true' do
-        expect(subject.enabled?).to be_truthy
+        expect(subject.enabled?).to eq(true)
       end
     end
 
@@ -28,7 +28,7 @@ RSpec.describe Gitlab::RepositorySizeChecker do
       let(:limit) { 0 }
 
       it 'returns false' do
-        expect(subject.enabled?).to be_falsey
+        expect(subject.enabled?).to eq(false)
       end
     end
   end
@@ -37,11 +37,11 @@ RSpec.describe Gitlab::RepositorySizeChecker do
     let(:current_size) { 49 }
 
     it 'returns true when changes go over' do
-      expect(subject.changes_will_exceed_size_limit?(2)).to be_truthy
+      expect(subject.changes_will_exceed_size_limit?(2.megabytes)).to eq(true)
     end
 
     it 'returns false when changes do not go over' do
-      expect(subject.changes_will_exceed_size_limit?(1)).to be_falsey
+      expect(subject.changes_will_exceed_size_limit?(1.megabytes)).to eq(false)
     end
   end
 

spec/lib/gitlab/repository_size_error_message_spec.rb

@@ -3,11 +3,11 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::RepositorySizeErrorMessage do
+  let_it_be(:namespace) { build(:namespace) }
   let(:checker) do
     Gitlab::RepositorySizeChecker.new(
       current_size_proc: -> { 15.megabytes },
-      total_repository_size_excess: 0,
-      additional_purchased_storage: 0,
+      namespace: namespace,
       limit: 10.megabytes
     )
   end
@@ -15,6 +15,10 @@ RSpec.describe Gitlab::RepositorySizeErrorMessage do
   let(:message) { checker.error_message }
   let(:base_message) { 'because this repository has exceeded its size limit of 10 MB by 5 MB' }
 
+  before do
+    allow(namespace).to receive(:total_repository_size_excess).and_return(0)
+  end
+
   describe 'error messages' do
     describe '#commit_error' do
       it 'returns the correct message' do

spec/models/snippet_spec.rb

@@ -666,8 +666,7 @@ RSpec.describe Snippet do
 
     let(:checker) { subject.repository_size_checker }
     let(:current_size) { 60 }
-    let(:total_repository_size_excess) { 0 }
-    let(:additional_purchased_storage) { 0 }
+    let(:namespace) { nil }
 
     before do
       allow(subject.repository).to receive(:size).and_return(current_size)

spec/support/shared_examples/lib/gitlab/repository_size_checker_shared_examples.rb

@@ -29,7 +29,7 @@ RSpec.shared_examples 'checker size exceeded' do
     let(:current_size) { 51 }
 
     it 'returns zero' do
-      expect(subject.exceeded_size).to eq(1)
+      expect(subject.exceeded_size).to eq(1.megabytes)
     end
   end
 
@@ -37,7 +37,7 @@ RSpec.shared_examples 'checker size exceeded' do
     let(:current_size) { 50 }
 
     it 'returns zero' do
-      expect(subject.exceeded_size(1)).to eq(1)
+      expect(subject.exceeded_size(1.megabytes)).to eq(1.megabytes)
     end
   end
 
@@ -45,7 +45,7 @@ RSpec.shared_examples 'checker size exceeded' do
     let(:current_size) { 49 }
 
     it 'returns zero' do
-      expect(subject.exceeded_size(1)).to eq(0)
+      expect(subject.exceeded_size(1.megabytes)).to eq(0)
     end
   end
 end

spec/support/shared_examples/models/snippet_shared_examples.rb

@@ -4,8 +4,7 @@ RSpec.shared_examples 'size checker for snippet' do |action|
   it 'sets up size checker', :aggregate_failures do
     expect(checker.current_size).to eq(current_size.megabytes)
     expect(checker.limit).to eq(Gitlab::CurrentSettings.snippet_size_limit)
-    expect(checker.total_repository_size_excess).to eq(total_repository_size_excess)
-    expect(checker.additional_purchased_storage).to eq(additional_purchased_storage)
     expect(checker.enabled?).to eq(true)
+    expect(checker.instance_variable_get(:@namespace)).to eq(namespace)
   end
 end