1 files changed, 18 insertions, 22 deletions
diff --git a/spec/models/commit_status_spec.rb b/spec/models/commit_status_spec.rb
index 9ce9f0e13b5..e9257b08bca 100644
--- a/spec/models/commit_status_spec.rb
+++ b/spec/models/commit_status_spec.rb
@@ -44,6 +44,24 @@ RSpec.describe CommitStatus, feature_category: :continuous_integration do
   it { is_expected.not_to be_retried }
   it { expect(described_class.primary_key).to eq('id') }
 
+  describe '.switch_table_names' do
+    before do
+      stub_env('USE_CI_BUILDS_ROUTING_TABLE', flag_value)
+    end
+
+    context 'with the env flag disabled' do
+      let(:flag_value) { 'false' }
+
+      it { expect(described_class.switch_table_names).to eq(:ci_builds) }
+    end
+
+    context 'with the env flag enabled' do
+      let(:flag_value) { 'true' }
+
+      it { expect(described_class.switch_table_names).to eq(:p_ci_builds) }
+    end
+  end
+
   describe '#author' do
     subject { commit_status.author }
 
@@ -1067,26 +1085,4 @@ RSpec.describe CommitStatus, feature_category: :continuous_integration do
 
     it_behaves_like 'having enum with nil value'
   end
-
-  describe 'routing table switch' do
-    context 'with ff disabled' do
-      before do
-        stub_feature_flags(ci_partitioning_use_ci_builds_routing_table: false)
-      end
-
-      it 'uses the legacy table' do
-        expect(described_class.table_name).to eq('ci_builds')
-      end
-    end
-
-    context 'with ff enabled' do
-      before do
-        stub_feature_flags(ci_partitioning_use_ci_builds_routing_table: true)
-      end
-
-      it 'uses the routing table' do
-        expect(described_class.table_name).to eq('p_ci_builds')
-      end
-    end
-  end
 end