Introduce "stub_env" test helper for safely stubbing environment variables

3 files changed, 14 insertions, 5 deletions

diff --git a/spec/db/production/settings.rb b/spec/db/production/settings.rb
index a7c5283df94..007b35bbb77 100644
--- a/spec/db/production/settings.rb
+++ b/spec/db/production/settings.rb
@@ -2,10 +2,11 @@ require 'spec_helper'
 require 'rainbow/ext/string'
 
 describe 'seed production settings', lib: true do
+  include StubENV
+
   context 'GITLAB_SHARED_RUNNERS_REGISTRATION_TOKEN is set in the environment' do
     before do
-      allow(ENV).to receive(:[]).and_call_original
-      allow(ENV).to receive(:[]).with('GITLAB_SHARED_RUNNERS_REGISTRATION_TOKEN').and_return('013456789')
+      stub_env('GITLAB_SHARED_RUNNERS_REGISTRATION_TOKEN', '013456789')
     end
 
     it 'writes the token to the database' do
diff --git a/spec/initializers/secret_token_spec.rb b/spec/initializers/secret_token_spec.rb
index 837b0de9a4c..ad7f032d1e5 100644
--- a/spec/initializers/secret_token_spec.rb
+++ b/spec/initializers/secret_token_spec.rb
@@ -2,10 +2,11 @@ require 'spec_helper'
 require_relative '../../config/initializers/secret_token'
 
 describe 'create_tokens', lib: true do
+  include StubENV
+
   let(:secrets) { ActiveSupport::OrderedOptions.new }
 
   before do
-    allow(ENV).to receive(:[]).and_call_original
     allow(File).to receive(:write)
     allow(File).to receive(:delete)
     allow(Rails).to receive_message_chain(:application, :secrets).and_return(secrets)
@@ -17,7 +18,7 @@ describe 'create_tokens', lib: true do
   context 'setting secret_key_base and otp_key_base' do
     context 'when none of the secrets exist' do
       before do
-        allow(ENV).to receive(:[]).with('SECRET_KEY_BASE').and_return(nil)
+        stub_env('SECRET_KEY_BASE', nil)
         allow(File).to receive(:exist?).with('.secret').and_return(false)
         allow(File).to receive(:exist?).with('config/secrets.yml').and_return(false)
         allow(self).to receive(:warn_missing_secret)
@@ -69,7 +70,7 @@ describe 'create_tokens', lib: true do
 
       context 'when secret_key_base exists in the environment and secrets.yml' do
         before do
-          allow(ENV).to receive(:[]).with('SECRET_KEY_BASE').and_return('env_key')
+          stub_env('SECRET_KEY_BASE', 'env_key')
           secrets.secret_key_base = 'secret_key_base'
           secrets.otp_key_base = 'otp_key_base'
         end
diff --git a/spec/support/stub_env.rb b/spec/support/stub_env.rb
new file mode 100644
index 00000000000..18597b5c71f
--- /dev/null
+++ b/spec/support/stub_env.rb
@@ -0,0 +1,7 @@
+module StubENV
+  def stub_env(key, value)
+    allow(ENV).to receive(:[]).and_call_original unless @env_already_stubbed
+    @env_already_stubbed ||= true
+    allow(ENV).to receive(:[]).with(key).and_return(value)
+  end
+end