diff --git a/lib/puppet/util/execution.rb b/lib/puppet/util/execution.rb index dd820f856..69f4f2c15 100644 --- a/lib/puppet/util/execution.rb +++ b/lib/puppet/util/execution.rb @@ -1,21 +1,20 @@ module Puppet::Util::Execution module_function # Run some code with a specific environment. Resets the environment back to # what it was at the end of the code. def withenv(hash) - oldvals = {} + saved = ENV.to_hash hash.each do |name, val| - name = name.to_s - oldvals[name] = ENV[name] - ENV[name] = val + ENV[name.to_s] = val end yield ensure - oldvals.each do |name, val| + ENV.clear + saved.each do |name, val| ENV[name] = val end end end diff --git a/spec/unit/util/execution_spec.rb b/spec/unit/util/execution_spec.rb new file mode 100644 index 000000000..312dd3b8e --- /dev/null +++ b/spec/unit/util/execution_spec.rb @@ -0,0 +1,49 @@ +#!/usr/bin/env ruby + +require File.dirname(__FILE__) + '/../../spec_helper' + +describe Puppet::Util::Execution do + include Puppet::Util::Execution + describe "#withenv" do + before :each do + @original_path = ENV["PATH"] + @new_env = {:PATH => "/some/bogus/path"} + end + + it "should change environment variables within the block then reset environment variables to their original values" do + withenv @new_env do + ENV["PATH"].should == "/some/bogus/path" + end + ENV["PATH"].should == @original_path + end + + it "should reset environment variables to their original values even if the block fails" do + begin + withenv @new_env do + ENV["PATH"].should == "/some/bogus/path" + raise "This is a failure" + end + rescue + end + ENV["PATH"].should == @original_path + end + + it "should reset environment variables even when they are set twice" do + # Setting Path & Environment parameters in Exec type can cause weirdness + @new_env["PATH"] = "/someother/bogus/path" + withenv @new_env do + # When assigning duplicate keys, can't guarantee order of evaluation + ENV["PATH"].should =~ /\/some.*\/bogus\/path/ + end + ENV["PATH"].should == @original_path + end + + it "should remove any new environment variables after the block ends" do + @new_env[:FOO] = "bar" + withenv @new_env do + ENV["FOO"].should == "bar" + end + ENV["FOO"].should == nil + end + end +end