From 19be5f141eb854941ebf71674998f7aa88132ad7 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 17 May 2010 14:12:11 -0700
Subject: [PATCH] `vagrant package` works with multi-VM environments

---
 lib/vagrant/commands/package.rb       |  69 ++++++++++---
 templates/strings.yml                 |   5 +
 test/vagrant/commands/package_test.rb | 141 +++++++++++++++-----------
 3 files changed, 141 insertions(+), 74 deletions(-)

diff --git a/lib/vagrant/commands/package.rb b/lib/vagrant/commands/package.rb
index 7b6ca6295..197342575 100644
--- a/lib/vagrant/commands/package.rb
+++ b/lib/vagrant/commands/package.rb
@@ -8,23 +8,58 @@ module Vagrant
       description "Packages a vagrant environment for distribution"
 
       def execute(args=[])
-        parse_options(args)
+        args = parse_options(args)
 
-        if !options[:base]
-          # Packaging a pre-existing environment
-          env.require_persisted_vm
+        if options[:base]
+          package_base
         else
-          # Packaging a base box; that is a VM not tied to a specific
-          # vagrant environment
-          vm = VM.find(options[:base])
-          vm.env = env if vm
-          env.vm = vm
+          package_single(args[0])
+        end
+      end
 
-          error_and_exit(:vm_base_not_found, :name => options[:base]) unless vm
+      def package_base
+        # Packaging a base box; that is a VM not tied to a specific
+        # vagrant environment
+        vm = VM.find(options[:base])
+        if !vm
+          error_and_exit(:vm_base_not_found, :name => options[:base])
+          return # for tests
         end
 
-        error_and_exit(:vm_power_off_to_package) unless env.vm.powered_off?
-        env.vm.package(args[0], options[:include])
+        vm.env = env
+        package_vm(vm)
+      end
+
+      def package_single(name)
+        if name.nil? && env.multivm?
+          error_and_exit(:package_multivm)
+          return
+        end
+
+        vm = if name.nil?
+               env.vms.values.first
+             else
+               env.vms[name.to_sym]
+             end
+
+        if vm.nil?
+          error_and_exit(:unknown_vm, :vm => name)
+          return
+        elsif !vm.created?
+          error_and_exit(:environment_not_created)
+          return
+        end
+
+        package_vm(vm)
+      end
+
+      def package_vm(vm)
+        if !vm.powered_off?
+          error_and_exit(:vm_power_off_to_package)
+          return # for tests
+        end
+
+        vm.package(options[:output], options[:include])
       end
 
       def options_spec(opts)
@@ -32,15 +67,21 @@ module Vagrant
 
         # Defaults
         options[:include] = []
+        options[:base] = nil
+        options[:output] = nil
 
-        opts.on("--base [BASE]", "Name or UUID of VM to create a base box from") do |v|
+        opts.on("--base BASE", "Name or UUID of VM to create a base box from") do |v|
           options[:base] = v
         end
 
         opts.on("--include x,y,z", Array, "List of files to include in the package") do |v|
           options[:include] = v
         end
+
+        opts.on("-o", "--output FILE", "File to save the package as.") do |v|
+          options[:output] = v
+        end
       end
     end
   end
-end
\ No newline at end of file
+end
diff --git a/templates/strings.yml b/templates/strings.yml
index 6be52ea14..1e5f87c65 100644
--- a/templates/strings.yml
+++ b/templates/strings.yml
@@ -108,6 +108,11 @@
   <%= Vagrant::Environment::ROOTFILE_NAME %> and running `vagrant up`
 :package_include_file_doesnt_exist: |-
   File specified to include: '<%= filename %>' does not exist!
+:package_multivm: |-
+  Because this Vagrant environment represents multiple VMs, a
+  specific VM must be specified. This can be done by calling
+  `vagrant package NAME` where NAME is a valid VM represented by
+  your Vagrantfile.
 :package_requires_export: |-
   Package must be used in conjunction with export.
 :provisioner_invalid_class: |-
diff --git a/test/vagrant/commands/package_test.rb b/test/vagrant/commands/package_test.rb
index 064d00a02..14d035670 100644
--- a/test/vagrant/commands/package_test.rb
+++ b/test/vagrant/commands/package_test.rb
@@ -4,81 +4,102 @@ class CommandsPackageTest < Test::Unit::TestCase
   setup do
     @klass = Vagrant::Commands::Package
 
-    @persisted_vm = mock("persisted_vm")
-    @persisted_vm.stubs(:execute!)
-
     @env = mock_environment
-    @env.stubs(:require_persisted_vm)
-    @env.stubs(:vm).returns(@persisted_vm)
-
     @instance = @klass.new(@env)
   end
 
   context "executing" do
+    should "package base if a base is given" do
+      @instance.expects(:package_base).once
+      @instance.execute(["--base","foo"])
+    end
+
+    should "package single if no name is given" do
+      @instance.expects(:package_single).with(nil).once
+      @instance.execute
+    end
+
+    should "package single if a name is given" do
+      @instance.expects(:package_single).with("foo").once
+      @instance.execute(["foo"])
+    end
+  end
+
+  context "packaging base" do
+    should "error and exit if no VM is found" do
+      Vagrant::VM.expects(:find).with("foo").returns(nil)
+      @instance.expects(:error_and_exit).with(:vm_base_not_found, :name => "foo").once
+      @instance.execute(["--base", "foo"])
+    end
+
+    should "package the VM like any other VM" do
+      vm = mock("vm")
+      Vagrant::VM.expects(:find).with("foo").returns(vm)
+      vm.expects(:env=).with(@env).once
+      @instance.expects(:package_vm).with(vm).once
+      @instance.execute(["--base", "foo"])
+    end
+  end
+
+  context "packaging a single VM" do
     setup do
-      @persisted_vm.stubs(:package)
-      @persisted_vm.stubs(:powered_off?).returns(true)
+      @vm = mock("vm")
+      @vm.stubs(:created?).returns(true)
+
+      @vms = {:bar => @vm}
+      @env.stubs(:vms).returns(@vms)
+      @env.stubs(:multivm?).returns(false)
     end
 
-    context "with no base specified" do
-      should "require a persisted vm" do
-        @env.expects(:require_persisted_vm).once
-        @instance.execute
-      end
+    should "error and exit if no name is given in a multi-vm env" do
+      @env.stubs(:multivm?).returns(true)
+      @instance.expects(:error_and_exit).with(:package_multivm).once
+      @instance.package_single(nil)
     end
 
-    context "with base specified" do
-      setup do
-        @vm = mock("vm")
-
-        Vagrant::VM.stubs(:find).with(@name).returns(@vm)
-        @vm.stubs(:env=).with(@env)
-        @env.stubs(:vm=)
-
-        @name = "bar"
-      end
-
-      should "find the given base and set it on the env" do
-        Vagrant::VM.expects(:find).with(@name).returns(@vm)
-        @vm.expects(:env=).with(@env)
-        @env.expects(:vm=).with(@vm)
-
-        @instance.execute(["foo", "--base", @name])
-      end
-
-      should "error if the VM is not found" do
-        Vagrant::VM.expects(:find).with(@name).returns(nil)
-        @instance.expects(:error_and_exit).with(:vm_base_not_found, :name => @name).once
-
-        @instance.execute(["foo", "--base", @name])
-      end
+    should "error and exit if the VM doesn't exist" do
+      @instance.expects(:error_and_exit).with(:unknown_vm, :vm => :foo).once
+      @instance.package_single(:foo)
     end
 
-    context "shared (with and without base specified)" do
-      should "error and exit if the VM is not powered off" do
-        @persisted_vm.stubs(:powered_off?).returns(false)
-        @instance.expects(:error_and_exit).with(:vm_power_off_to_package).once
-        @persisted_vm.expects(:package).never
-        @instance.execute
-      end
+    should "error and exit if the VM is not created" do
+      @vm.stubs(:created?).returns(false)
+      @instance.expects(:error_and_exit).with(:environment_not_created).once
+      @instance.package_single(:bar)
+    end
 
-      should "call package on the persisted VM" do
-        @persisted_vm.expects(:package).once
-        @instance.execute
-      end
+    should "use the first VM is no name is given in a single VM environment" do
+      @instance.expects(:package_vm).with(@vm).once
+      @instance.package_single(nil)
+    end
 
-      should "pass the out path and include_files to the package method" do
-        out_path = mock("out_path")
-        include_files = "foo"
-        @persisted_vm.expects(:package).with(out_path, [include_files]).once
-        @instance.execute([out_path, "--include", include_files])
-      end
+    should "package the VM" do
+      @instance.expects(:package_vm).with(@vm).once
+      @instance.package_single(:bar)
+    end
+  end
 
-      should "default to an empty array when not include_files are specified" do
-        out_path = mock("out_path")
-        @persisted_vm.expects(:package).with(out_path, []).once
-        @instance.execute([out_path])
-      end
+  context "packaging a VM" do
+    setup do
+      @vm = mock("vm")
+      @vm.stubs(:powered_off?).returns(true)
+
+      @options = {}
+      @instance.stubs(:options).returns(@options)
+    end
+
+    should "error and exit if VM is not powered off" do
+      @vm.stubs(:powered_off?).returns(false)
+      @instance.expects(:error_and_exit).with(:vm_power_off_to_package).once
+      @instance.package_vm(@vm)
+    end
+
+    should "package the VM with the proper arguments" do
+      @options[:output] = "foo.box"
+      @options[:include] = :bar
+
+      @vm.expects(:package).with(@options[:output], @options[:include]).once
+      @instance.package_vm(@vm)
     end
   end
 end