When merging orders, set distributor when target order distributor is nil

app/controllers/spree/base_controller_decorator.rb

@@ -12,6 +12,7 @@ module Spree
             session[:order_id] = last_incomplete_order.id
           elsif current_order && last_incomplete_order && current_order != last_incomplete_order
             if current_order.distributor.nil? || current_order.distributor == last_incomplete_order.distributor
+              current_order.set_distributor! last_incomplete_order.distributor if current_order.distributor.nil?
               current_order.merge!(last_incomplete_order)
             else
               last_incomplete_order.destroy

app/controllers/spree/orders_controller_decorator.rb

@@ -7,8 +7,7 @@ Spree::OrdersController.class_eval do
 
     if populate_valid? @distributor
       order = current_order(true)
-      order.distributor = @distributor
-      order.save!
+      order.set_distributor! @distributor
 
     else
       flash[:error] = "Please choose a distributor for this order." if @distributor.nil?

app/models/spree/order_decorator.rb

@@ -4,7 +4,6 @@ Spree::Order.class_eval do
   before_validation :shipping_address_from_distributor
   after_create :set_default_shipping_method
 
-
   def can_change_distributor?
     # Distributor may not be changed once an item has been added to the cart/order
     line_items.empty?
@@ -15,6 +14,12 @@ Spree::Order.class_eval do
     super(distributor)
   end
 
+  def set_distributor!(distributor)
+    self.distributor = distributor
+    save!
+  end
+
+
   def can_add_product_to_cart?(product)
     can_change_distributor? || product.distributors.include?(distributor)
   end

spec/controllers/home_controller_spec.rb

@@ -43,5 +43,21 @@ describe Spree::HomeController do
 
       spree_get :index
     end
+
+    it "sets the distributor when the target order has no distributor" do
+      incomplete_order = double(:order, distributor: 1)
+      current_order = double(:order, distributor: nil)
+
+      user = double(:user, last_incomplete_order: incomplete_order)
+      controller.stub(:current_user).and_return(user)
+      controller.stub(:current_order).and_return(current_order)
+
+      current_order.should_receive(:set_distributor!).with(1)
+      current_order.should_receive(:merge!)
+
+      session[:order_id] = 123
+
+      spree_get :index
+    end
   end
 end