1 files changed, 13 insertions, 12 deletions
diff --git a/spec/lib/marginalia_spec.rb b/spec/lib/marginalia_spec.rb
index a920f598c24..fa0cd214c7e 100644
--- a/spec/lib/marginalia_spec.rb
+++ b/spec/lib/marginalia_spec.rb
@@ -24,18 +24,6 @@ RSpec.describe 'Marginalia spec' do
     end
   end
 
-  def stub_feature(value)
-    allow(Gitlab::Marginalia).to receive(:cached_feature_enabled?).and_return(value)
-  end
-
-  def make_request(correlation_id)
-    request_env = Rack::MockRequest.env_for('/')
-
-    ::Labkit::Correlation::CorrelationId.use_id(correlation_id) do
-      MarginaliaTestController.action(:first_user).call(request_env)
-    end
-  end
-
   describe 'For rails web requests' do
     let(:correlation_id) { SecureRandom.uuid }
     let(:recorded) { ActiveRecord::QueryRecorder.new { make_request(correlation_id) } }
@@ -149,4 +137,17 @@ RSpec.describe 'Marginalia spec' do
       end
     end
   end
+
+  def stub_feature(value)
+    stub_feature_flags(marginalia: value)
+    Gitlab::Marginalia.set_enabled_from_feature_flag
+  end
+
+  def make_request(correlation_id)
+    request_env = Rack::MockRequest.env_for('/')
+
+    ::Labkit::Correlation::CorrelationId.use_id(correlation_id) do
+      MarginaliaTestController.action(:first_user).call(request_env)
+    end
+  end
 end