Move state_machine's additions inside class_eval

And also cover them with tests.

app/models/spree/order_decorator.rb

@@ -45,6 +45,12 @@ Spree::Order.class_eval do
     remove_transition :from => :delivery, :to => :confirm
   end
 
+  state_machine.after_transition to: :payment, do: :charge_shipping_and_payment_fees!
+
+  state_machine.event :restart_checkout do
+    transition to: :cart, unless: :completed?
+  end
+
   # -- Scopes
   scope :managed_by, lambda { |user|
     if user.has_spree_role?('admin')
@@ -414,17 +420,13 @@ Spree::Order.class_eval do
     adjustment.state = state
   end
 
-  # object_params sets the payment amount to the order total, but it does this before
-  # the shipping method is set. This results in the customer not being charged for their
-  # order's shipping. To fix this, we refresh the payment amount here.
+  # object_params sets the payment amount to the order total, but it does this
+  # before the shipping method is set. This results in the customer not being
+  # charged for their order's shipping. To fix this, we refresh the payment
+  # amount here.
   def charge_shipping_and_payment_fees!
     update_totals
     return unless payments.any?
     payments.first.update_attribute :amount, total
   end
 end
-
-Spree::Order.state_machine.after_transition to: :payment, do: :charge_shipping_and_payment_fees!
-Spree::Order.state_machine.event :restart_checkout do
-  transition :to => :cart, unless: :completed?
-end

spec/models/spree/order_spec.rb

@@ -858,4 +858,46 @@ describe Spree::Order do
       end
     end
   end
+
+  describe '#restart_checkout!' do
+    let(:order) { build(:order) }
+
+    context 'when the order is complete' do
+      before { order.completed_at = Time.zone.now }
+
+      it 'raises' do
+        expect { order.restart_checkout! }
+          .to raise_error(StateMachine::InvalidTransition)
+      end
+    end
+
+    context 'when the is not complete' do
+      before { order.completed_at = nil }
+
+      it 'transitions to :cart state' do
+        order.restart_checkout!
+        expect(order.state).to eq('cart')
+      end
+    end
+  end
+
+  describe '#charge_shipping_and_payment_fees!' do
+    let(:order) do
+      build(:order, shipping_method: build(:shipping_method))
+    end
+
+    context 'after transitioning to payment' do
+      before do
+        order.state = 'delivery' # payment's previous state
+
+        allow(order).to receive(:payment_required?) { true }
+        allow(order).to receive(:charge_shipping_and_payment_fees!)
+      end
+
+      it 'calls charge_shipping_and_payment_fees!' do
+        order.next
+        expect(order).to have_received(:charge_shipping_and_payment_fees!)
+      end
+    end
+  end
 end