1 files changed, 21 insertions, 59 deletions
diff --git a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
index 3d60457c3a9..8f1f5b5ba1b 100644
--- a/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
+++ b/spec/lib/gitlab/database/partitioning/partition_manager_spec.rb
@@ -12,26 +12,18 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
   end
 
-  describe '.register' do
-    let(:model) { double(partitioning_strategy: nil) }
-
-    it 'remembers registered models' do
-      expect { described_class.register(model) }.to change { described_class.models }.to include(model)
-    end
-  end
-
   context 'creating partitions (mocked)' do
-    subject(:sync_partitions) { described_class.new(models).sync_partitions }
+    subject(:sync_partitions) { described_class.new(model).sync_partitions }
 
-    let(:models) { [model] }
-    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
+    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: []) }
+    let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "some_table" }
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
-      allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+      allow(connection).to receive(:table_exists?).and_call_original
+      allow(connection).to receive(:table_exists?).with(table).and_return(true)
+      allow(connection).to receive(:execute).and_call_original
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -44,35 +36,23 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
 
     it 'creates the partition' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
-      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
+      expect(connection).to receive(:execute).with(partitions.first.to_sql)
+      expect(connection).to receive(:execute).with(partitions.second.to_sql)
 
       sync_partitions
     end
 
-    context 'error handling with 2 models' do
-      let(:models) do
-        [
-          double(partitioning_strategy: strategy1, table_name: table),
-          double(partitioning_strategy: strategy2, table_name: table)
-        ]
-      end
-
-      let(:strategy1) { double('strategy1', missing_partitions: nil, extra_partitions: []) }
-      let(:strategy2) { double('strategy2', missing_partitions: partitions, extra_partitions: []) }
+    context 'when an error occurs during partition management' do
+      it 'does not raise an error' do
+        expect(partitioning_strategy).to receive(:missing_partitions).and_raise('this should never happen (tm)')
 
-      it 'still creates partitions for the second table' do
-        expect(strategy1).to receive(:missing_partitions).and_raise('this should never happen (tm)')
-        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
-        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
-
-        sync_partitions
+        expect { sync_partitions }.not_to raise_error
       end
     end
   end
 
   context 'creating partitions' do
-    subject(:sync_partitions) { described_class.new([my_model]).sync_partitions }
+    subject(:sync_partitions) { described_class.new(my_model).sync_partitions }
 
     let(:connection) { ActiveRecord::Base.connection }
     let(:my_model) do
@@ -101,15 +81,15 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
   context 'detaching partitions (mocked)' do
     subject(:sync_partitions) { manager.sync_partitions }
 
-    let(:manager) { described_class.new(models) }
-    let(:models) { [model] }
-    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table)}
+    let(:manager) { described_class.new(model) }
+    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: []) }
+    let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "foo" }
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
+      allow(connection).to receive(:table_exists?).and_call_original
+      allow(connection).to receive(:table_exists?).with(table).and_return(true)
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -131,24 +111,6 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
         sync_partitions
       end
-
-      context 'error handling' do
-        let(:models) do
-          [
-            double(partitioning_strategy: error_strategy, table_name: table),
-            model
-          ]
-        end
-
-        let(:error_strategy) { double(extra_partitions: nil, missing_partitions: []) }
-
-        it 'still drops partitions for the other model' do
-          expect(error_strategy).to receive(:extra_partitions).and_raise('injected error!')
-          extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) }
-
-          sync_partitions
-        end
-      end
     end
 
     context 'with the partition_pruning feature flag disabled' do
@@ -171,7 +133,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       end
     end
 
-    subject { described_class.new([my_model]).sync_partitions }
+    subject { described_class.new(my_model).sync_partitions }
 
     let(:connection) { ActiveRecord::Base.connection }
     let(:my_model) do
@@ -280,11 +242,11 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     it 'creates partitions for the future then drops the oldest one after a month' do
       # 1 month for the current month, 1 month for the old month that we're retaining data for, headroom
       expected_num_partitions = (Gitlab::Database::Partitioning::MonthlyStrategy::HEADROOM + 2.months) / 1.month
-      expect { described_class.new([my_model]).sync_partitions }.to change { num_partitions(my_model) }.from(0).to(expected_num_partitions)
+      expect { described_class.new(my_model).sync_partitions }.to change { num_partitions(my_model) }.from(0).to(expected_num_partitions)
 
       travel 1.month
 
-      expect { described_class.new([my_model]).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
+      expect { described_class.new(my_model).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
     end
   end
 end