use version of stub_env from main code

borrowed from `spec/support/stub_env.rb`

3 files changed, 49 insertions, 7 deletions

diff --git a/qa/qa.rb b/qa/qa.rb
index 40bdcc726ae..242f9851a7f 100644
--- a/qa/qa.rb
+++ b/qa/qa.rb
@@ -147,6 +147,14 @@ module QA
     autoload :Config, 'qa/specs/config'
     autoload :Runner, 'qa/specs/runner'
   end
+
+
+  ##
+  # Classes that provide support methods
+  #
+  module Support
+    autoload :StubENV, 'qa/support/stub_env'
+  end
 end
 
 QA::Runtime::Release.extend_autoloads!
diff --git a/qa/qa/support/stub_env.rb b/qa/qa/support/stub_env.rb
new file mode 100644
index 00000000000..74beb62151f
--- /dev/null
+++ b/qa/qa/support/stub_env.rb
@@ -0,0 +1,40 @@
+# Inspired by https://github.com/ljkbennett/stub_env/blob/master/lib/stub_env/helpers.rb
+module QA
+  module Support
+    module StubENV
+      def stub_env(key_or_hash, value = nil)
+        init_stub unless env_stubbed?
+
+        if key_or_hash.is_a? Hash
+          key_or_hash.each { |k, v| add_stubbed_value(k, v) }
+        else
+          add_stubbed_value key_or_hash, value
+        end
+      end
+
+      private
+
+      STUBBED_KEY = '__STUBBED__'.freeze
+
+      def add_stubbed_value(key, value)
+        allow(ENV).to receive(:[]).with(key).and_return(value)
+        allow(ENV).to receive(:key?).with(key).and_return(true)
+        allow(ENV).to receive(:fetch).with(key).and_return(value)
+        allow(ENV).to receive(:fetch).with(key, anything()) do |_, default_val|
+          value || default_val
+        end
+      end
+
+      def env_stubbed?
+        ENV[STUBBED_KEY]
+      end
+
+      def init_stub
+        allow(ENV).to receive(:[]).and_call_original
+        allow(ENV).to receive(:key?).and_call_original
+        allow(ENV).to receive(:fetch).and_call_original
+        add_stubbed_value(STUBBED_KEY, true)
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/qa/spec/runtime/env_spec.rb b/qa/spec/runtime/env_spec.rb
index 57a72a04507..3914a514244 100644
--- a/qa/spec/runtime/env_spec.rb
+++ b/qa/spec/runtime/env_spec.rb
@@ -1,7 +1,5 @@
 describe QA::Runtime::Env do
-  before do
-    allow(ENV).to receive(:[]).and_call_original
-  end
+  include QA::Support::StubENV
 
   describe '.chrome_headless?' do
     context 'when there is an env variable set' do
@@ -57,8 +55,4 @@ describe QA::Runtime::Env do
       end
     end
   end
-
-  def stub_env(name, value)
-    allow(ENV).to receive(:[]).with(name).and_return(value)
-  end
 end