1 files changed, 72 insertions, 1 deletions
diff --git a/spec/models/namespace/aggregation_schedule_spec.rb b/spec/models/namespace/aggregation_schedule_spec.rb
index 5ba7547ff4d..8ed0248e1b2 100644
--- a/spec/models/namespace/aggregation_schedule_spec.rb
+++ b/spec/models/namespace/aggregation_schedule_spec.rb
@@ -2,6 +2,77 @@
 
 require 'spec_helper'
 
-RSpec.describe Namespace::AggregationSchedule, type: :model do
+RSpec.describe Namespace::AggregationSchedule, :clean_gitlab_redis_shared_state, type: :model do
+  include ExclusiveLeaseHelpers
+
   it { is_expected.to belong_to :namespace }
+
+  describe '.delay_timeout' do
+    context 'when timeout is set on redis' do
+      it 'uses personalized timeout' do
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(described_class::REDIS_SHARED_KEY, 1.hour)
+        end
+
+        expect(described_class.delay_timeout).to eq(1.hour)
+      end
+    end
+
+    context 'when timeout is not set on redis' do
+      it 'uses default timeout' do
+        expect(described_class.delay_timeout).to eq(3.hours)
+      end
+    end
+  end
+
+  describe '#schedule_root_storage_statistics' do
+    let(:namespace) { create(:namespace) }
+    let(:aggregation_schedule) { namespace.build_aggregation_schedule }
+    let(:lease_key) { "namespace:namespaces_root_statistics:#{namespace.id}" }
+
+    context "when we can't obtain the lease" do
+      it 'does not schedule the workers' do
+        stub_exclusive_lease_taken(lease_key, timeout: described_class::DEFAULT_LEASE_TIMEOUT)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .not_to receive(:perform_async)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .not_to receive(:perform_in)
+
+        aggregation_schedule.save!
+      end
+    end
+
+    context 'when we can obtain the lease' do
+      it 'schedules a root storage statistics after create' do
+        stub_exclusive_lease(lease_key, timeout: described_class::DEFAULT_LEASE_TIMEOUT)
+
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_async).once
+
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_in).once
+          .with(described_class::DEFAULT_LEASE_TIMEOUT, aggregation_schedule.namespace_id )
+
+        aggregation_schedule.save!
+      end
+    end
+
+    context 'with a personalized lease timeout' do
+      before do
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set(described_class::REDIS_SHARED_KEY, 1.hour)
+        end
+      end
+
+      it 'uses a personalized time' do
+        expect(Namespaces::RootStatisticsWorker)
+          .to receive(:perform_in)
+          .with(1.hour, aggregation_schedule.namespace_id)
+
+        aggregation_schedule.save!
+      end
+    end
+  end
 end