Adding a "withenv" execution util method, and using it in :exec for path handling. Next will be other env handling.

git-svn-id: https://reductivelabs.com/svn/puppet/trunk@1567 980ebf18-57e1-0310-9a29-db15c13687c0

4 files changed, 95 insertions, 36 deletions

diff --git a/lib/puppet/type/exec.rb b/lib/puppet/type/exec.rb
index 0dc3ec02a..c5cf81142 100755
--- a/lib/puppet/type/exec.rb
+++ b/lib/puppet/type/exec.rb
@@ -1,5 +1,7 @@
 module Puppet
     newtype(:exec) do
+        include Puppet::Util::Execution
+
         @doc = "Executes external commands.  It is critical that all commands
             executed using this mechanism can be run multiple times without
             harm, i.e., they are *idempotent*.  One useful way to create idempotent
@@ -61,6 +63,7 @@ module Puppet
         end
 
         newstate(:returns) do |state|
+            include Puppet::Util::Execution
             munge do |value|
                 value.to_s
             end
@@ -93,16 +96,14 @@ module Puppet
                     end
                 elsif path = self.parent[:path]
                     exe = cmd.split(/ /)[0]
-                    tmppath = ENV["PATH"]
-                    ENV["PATH"] = self.parent[:path].join(":")
-
-                    path = %{which #{exe}}.chomp
-                    if path == ""
-                        self.fail(
-                            "Could not find command '%s'" % exe
-                        )
+                    withenv :PATH => self.parent[:path].join(":") do
+                        path = %{which #{exe}}.chomp
+                        if path == ""
+                            self.fail(
+                                "Could not find command '%s'" % exe
+                            )
+                        end
                     end
-                    ENV["PATH"] = tmppath
                 else
                     self.fail(
                         "%s is somehow not qualified with no search path" %
@@ -130,7 +131,6 @@ module Puppet
 
                 # We need a dir to change to, even if it's just the cwd
                 dir = self.parent[:cwd] || Dir.pwd
-                tmppath = ENV["PATH"]
 
                 event = :executed_command
 
@@ -470,7 +470,6 @@ module Puppet
         def run(command, check = false)
             output = nil
             status = nil
-            tmppath = ENV["PATH"]
 
             dir = nil
 
@@ -493,39 +492,40 @@ module Puppet
             end
             begin
                 # Do our chdir
-                Dir.chdir(dir) {
+                Dir.chdir(dir) do
+                    env = {}
+
                     if self[:path]
-                        ENV["PATH"] = self[:path].join(":")
+                        env[:PATH] = self[:path].join(":")
                     end
 
-                    # The user and group default to nil, which 'asuser'
-                    # handlers correctly
-                    Puppet::Util.asuser(self[:user], self[:group]) {
-                        # capture both stdout and stderr
-                        if self[:user]
-                            unless defined? @@alreadywarned
-                                Puppet.warning(
-                        "Cannot capture STDERR when running as another user"
-                                )
-                                @@alreadywarned = true
+                    withenv env do
+                        # The user and group default to nil, which 'asuser'
+                        # handlers correctly
+                        Puppet::Util.asuser(self[:user], self[:group]) {
+                            # capture both stdout and stderr
+                            if self[:user]
+                                unless defined? @@alreadywarned
+                                    Puppet.warning(
+                            "Cannot capture STDERR when running as another user"
+                                    )
+                                    @@alreadywarned = true
+                                end
+                                output = %x{#{command}}
+                            else
+                                output = %x{#{command} 2>&1}
                             end
-                            output = %x{#{command}}
-                        else
-                            output = %x{#{command} 2>&1}
-                        end
-                    }
-                    status = $?.dup
+                        }
+                        status = $?.dup
 
-                    # The shell returns 127 if the command is missing.
-                    if $?.exitstatus == 127
-                        raise ArgumentError, output
+                        # The shell returns 127 if the command is missing.
+                        if $?.exitstatus == 127
+                            raise ArgumentError, output
+                        end
                     end
-                }
+                end
             rescue Errno::ENOENT => detail
                 self.fail detail.to_s
-            ensure
-                # reset things to how we found them
-                ENV["PATH"] = tmppath
             end
 
             return output, status
diff --git a/lib/puppet/util.rb b/lib/puppet/util.rb
index a818a77bf..f42eddef0 100644
--- a/lib/puppet/util.rb
+++ b/lib/puppet/util.rb
@@ -520,5 +520,6 @@ require 'puppet/util/methodhelper'
 require 'puppet/util/metaid'
 require 'puppet/util/classgen'
 require 'puppet/util/docs'
+require 'puppet/util/execution'
 
 # $Id$
diff --git a/lib/puppet/util/execution.rb b/lib/puppet/util/execution.rb
new file mode 100644
index 000000000..67b5ed692
--- /dev/null
+++ b/lib/puppet/util/execution.rb
@@ -0,0 +1,21 @@
+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 = {}
+        hash.each do |name, val|
+            name = name.to_s
+            oldvals[name] = ENV[name]
+        end
+
+        yield
+    ensure
+        oldvals.each do |name, val|
+            ENV[name] = val
+        end
+    end
+end
+
+# $Id$
diff --git a/test/util/execution.rb b/test/util/execution.rb
new file mode 100755
index 000000000..4df606806
--- /dev/null
+++ b/test/util/execution.rb
@@ -0,0 +1,37 @@
+if __FILE__ == $0
+    $:.unshift '..'
+    $:.unshift '../../lib'
+    $puppetbase = ".."
+end
+
+require 'puppet'
+require 'puppettest'
+require 'test/unit'
+
+class TestPuppetUtilExecution < Test::Unit::TestCase
+    include TestPuppet
+
+    def test_withenv
+        ENV["testing"] = "yay"
+
+        assert_nothing_raised do
+            Puppet::Util::Execution.withenv :testing => "foo" do
+                $ran = true
+            end
+        end
+
+        assert_equal("yay", ENV["testing"])
+        assert_equal(true, $ran)
+
+        ENV["rah"] = "yay"
+        assert_raise(ArgumentError) do
+            Puppet::Util::Execution.withenv :testing => "foo" do
+                raise ArgumentError, "yay"
+            end
+        end
+
+        assert_equal("yay", ENV["rah"])
+    end
+end
+
+# $Id$