From 8aa4631e210511d8cfc4a1289001276cb839513d Mon Sep 17 00:00:00 2001
From: GeorgeThoppil <george.thoppil@shopify.com>
Date: Thu, 10 Feb 2022 20:34:11 -0500
Subject: [PATCH] Setting payment state for subscription order

---
 app/models/spree/order.rb                     |  3 +-
 spec/models/spree/order/state_machine_spec.rb | 58 +++++++++++++------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/app/models/spree/order.rb b/app/models/spree/order.rb
index 1e58900e02..e401dda2f6 100644
--- a/app/models/spree/order.rb
+++ b/app/models/spree/order.rb
@@ -669,11 +669,12 @@ module Spree
       shipments.each(&:cancel!)
 
       OrderMailer.cancel_email(id).deliver_later if send_cancellation_email
-      self.payment_state = 'credit_owed' unless shipped?
+      update(payment_state: updater.update_payment_state)
     end
 
     def after_resume
       shipments.each(&:resume!)
+      update(payment_state: updater.update_payment_state)
     end
 
     def use_billing?
diff --git a/spec/models/spree/order/state_machine_spec.rb b/spec/models/spree/order/state_machine_spec.rb
index 8da1df767a..dc4b767ce7 100644
--- a/spec/models/spree/order/state_machine_spec.rb
+++ b/spec/models/spree/order/state_machine_spec.rb
@@ -136,20 +136,11 @@ describe Spree::Order do
       end
 
       context "without shipped items" do
-        it "should set payment state to 'credit owed'" do
-          order.cancel!
-          expect(order.payment_state).to eq 'credit_owed'
-        end
-      end
-
-      context "with shipped items" do
-        before do
-          allow(order).to receive_messages shipment_state: 'partial'
-        end
-
-        it "should not alter the payment state" do
-          order.cancel!
-          expect(order.payment_state).to be_nil
+        it "should set payment state to 'void'" do
+          expect {
+            order.cancel!
+            order.reload
+          }.to change { order.payment_state }.to("void")
         end
       end
     end
@@ -157,10 +148,41 @@ describe Spree::Order do
 
   # Another regression test for Spree #729
   context "#resume" do
-    before do
-      allow(order).to receive_messages email: "user@spreecommerce.com"
-      allow(order).to receive_messages state: "canceled"
-      allow(order).to receive_messages allow_resume?: true
+    context "resets payment state" do
+      let!(:variant) { build(:variant) }
+      before do
+        allow(order).to receive_messages email: "user@spreecommerce.com"
+        allow(order).to receive_messages allow_cancel?: true
+        allow(order).to receive_messages allow_resume?: true
+        allow(order).to receive_messages line_items:
+                                           [build(:line_item, variant: variant, quantity: 2)]
+        allow(order.line_items).to receive_messages find_by_variant_id: order.line_items.first
+        order.update(total: 10)
+        order.cancel!
+      end
+
+      it "should set payment state to 'balance due'" do
+        expect {
+          order.resume!
+          order.reload
+        }.to change { order.payment_state }.to("balance_due")
+      end
+
+      it "should set payment state to 'paid'" do
+        expect {
+          order.update(payment_total: 10)
+          order.resume!
+          order.reload
+        }.to change { order.payment_state }.to("paid")
+      end
+
+      it "should set payment state to 'credit owed'" do
+        expect {
+          order.update(payment_total: 20)
+          order.resume!
+          order.reload
+        }.to change { order.payment_state }.to("credit_owed")
+      end
     end
   end
 end