From 5dd3e51788021bab23f03c714e1fedaea21de6dd Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 28 Feb 2010 19:56:50 -0800
Subject: [PATCH] Export action now cleans up and rescues. Package no longer
 cleans up temporary directory.

---
 lib/vagrant/actions/vm/export.rb        | 11 +++++++++++
 lib/vagrant/actions/vm/package.rb       |  8 +-------
 test/vagrant/actions/vm/export_test.rb  | 25 +++++++++++++++++++++++++
 test/vagrant/actions/vm/package_test.rb | 16 +---------------
 4 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/lib/vagrant/actions/vm/export.rb b/lib/vagrant/actions/vm/export.rb
index dac743db0..d1749a7e4 100644
--- a/lib/vagrant/actions/vm/export.rb
+++ b/lib/vagrant/actions/vm/export.rb
@@ -9,6 +9,17 @@ module Vagrant
           export
         end
 
+        def cleanup
+          if temp_dir
+            logger.info "Removing temporary export directory..."
+            FileUtils.rm_r(temp_dir)
+          end
+        end
+
+        def rescue(exception)
+          cleanup
+        end
+
         def setup_temp_dir
           @temp_dir = File.join(Env.tmp_path, Time.now.to_i.to_s)
 
diff --git a/lib/vagrant/actions/vm/package.rb b/lib/vagrant/actions/vm/package.rb
index 5cbb5fae6..b825bcf6b 100644
--- a/lib/vagrant/actions/vm/package.rb
+++ b/lib/vagrant/actions/vm/package.rb
@@ -13,7 +13,7 @@ module Vagrant
           @temp_path = nil
         end
 
-        def prepare 
+        def prepare
           @include_files.each do |file|
             raise ActionException.new("#{file} does not exist") unless File.exists?(file)
           end
@@ -21,12 +21,6 @@ module Vagrant
 
         def execute!
           compress
-          clean
-        end
-
-        def clean
-          logger.info "Removing temporary directory ..."
-          FileUtils.rm_r(temp_path)
         end
 
         def tar_path
diff --git a/test/vagrant/actions/vm/export_test.rb b/test/vagrant/actions/vm/export_test.rb
index 182f48aa2..f79cf7d5d 100644
--- a/test/vagrant/actions/vm/export_test.rb
+++ b/test/vagrant/actions/vm/export_test.rb
@@ -65,4 +65,29 @@ class ExportActionTest < Test::Unit::TestCase
       @action.export
     end
   end
+
+  context "cleanup" do
+    setup do
+      @temp_dir = "foo"
+      @action.stubs(:temp_dir).returns(@temp_dir)
+    end
+
+    should "remove the temporary directory" do
+      FileUtils.expects(:rm_r).with(@temp_dir).once
+      @action.cleanup
+    end
+
+    should "not remove a directory if temp_dir is nil" do
+      FileUtils.expects(:rm_r).never
+      @action.stubs(:temp_dir).returns(nil)
+      @action.cleanup
+    end
+  end
+
+  context "rescue" do
+    should "call cleanup method" do
+      @action.expects(:cleanup).once
+      @action.rescue(nil)
+    end
+  end
 end
diff --git a/test/vagrant/actions/vm/package_test.rb b/test/vagrant/actions/vm/package_test.rb
index 1af45bfc3..53642c2b4 100644
--- a/test/vagrant/actions/vm/package_test.rb
+++ b/test/vagrant/actions/vm/package_test.rb
@@ -29,29 +29,15 @@ class PackageActionTest < Test::Unit::TestCase
   context "executing" do
     setup do
       @action.stubs(:compress)
-      @action.stubs(:clean)
     end
 
-    should "compress and remove the working directory" do
+    should "compress" do
       package_seq = sequence("package_seq")
       @action.expects(:compress).in_sequence(package_seq)
-      @action.expects(:clean).in_sequence(package_seq)
       @action.execute!
     end
   end
 
-  context "cleaning up" do
-    setup do
-      @working_dir = "foo"
-      @action.stubs(:working_dir).returns(@working_dir)
-    end
-
-    should "remove the working directory" do
-      FileUtils.expects(:rm_r).with(@temp_path).once
-      @action.clean
-    end
-  end
-
   context "tar path" do
     should "be the temporary directory with the name and extension attached" do
       pwd = "foo"