2 files changed, 30 insertions, 0 deletions
diff --git a/spec/models/concerns/ci/partitionable/switch_spec.rb b/spec/models/concerns/ci/partitionable/switch_spec.rb
index 551ae111fa4..4aaf69bb585 100644
--- a/spec/models/concerns/ci/partitionable/switch_spec.rb
+++ b/spec/models/concerns/ci/partitionable/switch_spec.rb
@@ -38,6 +38,7 @@ RSpec.describe Ci::Partitionable::Switch, :aggregate_failures do
         id serial NOT NULL PRIMARY KEY,
         job_id int,
         partition_id int NOT NULL DEFAULT 1,
+        type text,
         expanded_environment_name text);
 
       CREATE TABLE _test_p_ci_jobs_metadata (
@@ -89,6 +90,25 @@ RSpec.describe Ci::Partitionable::Switch, :aggregate_failures do
 
   it { expect(partitioned_model.sequence_name).to eq('_test_ci_jobs_metadata_id_seq') }
 
+  context 'with singe table inheritance' do
+    let(:child_model) do
+      Class.new(model) do
+        def self.name
+          'TestSwitchJobMetadataChild'
+        end
+      end
+    end
+
+    it 'adds a Partitioned model for each descendant' do
+      expect(model::Partitioned).not_to eq(child_model::Partitioned)
+    end
+
+    it 'uses the parent name in STI queries' do
+      recorder = ActiveRecord::QueryRecorder.new { child_model.all.load }
+      expect(recorder.log).to include(/"type" = 'TestSwitchJobMetadataChild'/)
+    end
+  end
+
   context 'when switching the tables' do
     before do
       stub_feature_flags(table_rollout_flag => false)
diff --git a/spec/models/concerns/ci/partitionable_spec.rb b/spec/models/concerns/ci/partitionable_spec.rb
index d41654e547e..6daafc78cff 100644
--- a/spec/models/concerns/ci/partitionable_spec.rb
+++ b/spec/models/concerns/ci/partitionable_spec.rb
@@ -25,7 +25,11 @@ RSpec.describe Ci::Partitionable do
   end
 
   context 'with through options' do
+    let(:disable_partitionable_switch) { nil }
+
     before do
+      stub_env('DISABLE_PARTITIONABLE_SWITCH', disable_partitionable_switch)
+
       allow(ActiveSupport::DescendantsTracker).to receive(:store_inherited)
       stub_const("#{described_class}::Testing::PARTITIONABLE_MODELS", [ci_model.name])
 
@@ -39,6 +43,12 @@ RSpec.describe Ci::Partitionable do
     it { expect(ci_model.routing_table_name_flag).to eq(:some_flag) }
 
     it { expect(ci_model.ancestors).to include(described_class::Switch) }
+
+    context 'when DISABLE_PARTITIONABLE_SWITCH is set' do
+      let(:disable_partitionable_switch) { true }
+
+      it { expect(ci_model.ancestors).not_to include(described_class::Switch) }
+    end
   end
 
   context 'with partitioned options' do