1 files changed, 13 insertions, 6 deletions
diff --git a/spec/lib/gitlab/experiment/rollout/feature_spec.rb b/spec/lib/gitlab/experiment/rollout/feature_spec.rb
index d73757be79b..82603e6fe0f 100644
--- a/spec/lib/gitlab/experiment/rollout/feature_spec.rb
+++ b/spec/lib/gitlab/experiment/rollout/feature_spec.rb
@@ -9,9 +9,10 @@ RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment do
 
   describe "#enabled?" do
     before do
-      allow(Feature::Definition).to receive(:get).and_return('_instance_')
-      allow(Gitlab).to receive(:dev_env_or_com?).and_return(true)
-      allow(Feature).to receive(:get).and_return(double(state: :on))
+      stub_feature_flags(gitlab_experiment: true)
+      allow(subject).to receive(:feature_flag_defined?).and_return(true)
+      allow(Gitlab).to receive(:com?).and_return(true)
+      allow(subject).to receive(:feature_flag_instance).and_return(double(state: :on))
     end
 
     it "is enabled when all criteria are met" do
@@ -19,19 +20,25 @@ RSpec.describe Gitlab::Experiment::Rollout::Feature, :experiment do
     end
 
     it "isn't enabled if the feature definition doesn't exist" do
-      expect(Feature::Definition).to receive(:get).with('namespaced_stub').and_return(nil)
+      expect(subject).to receive(:feature_flag_defined?).and_return(false)
 
       expect(subject).not_to be_enabled
     end
 
     it "isn't enabled if we're not in dev or dotcom environments" do
-      expect(Gitlab).to receive(:dev_env_or_com?).and_return(false)
+      expect(Gitlab).to receive(:com?).and_return(false)
 
       expect(subject).not_to be_enabled
     end
 
     it "isn't enabled if the feature flag state is :off" do
-      expect(Feature).to receive(:get).with('namespaced_stub').and_return(double(state: :off))
+      expect(subject).to receive(:feature_flag_instance).and_return(double(state: :off))
+
+      expect(subject).not_to be_enabled
+    end
+
+    it "isn't enabled if the gitlab_experiment feature flag is false" do
+      stub_feature_flags(gitlab_experiment: false)
 
       expect(subject).not_to be_enabled
     end