fix(wallet): Include draft invoices inside wallet's ongoing balance

app/models/customer.rb

@@ -199,6 +199,12 @@ def reset_dunning_campaign!
     )
   end
 
+  def flag_wallets_for_refresh
+    return unless wallets.active.any?
+
+    wallets.active.update_all(ready_to_be_refreshed: true) # rubocop:disable Rails/SkipsModelValidations
+  end
+
   private
 
   def ensure_slug

app/models/wallet.rb

@@ -10,8 +10,9 @@ class Wallet < ApplicationRecord
   has_many :wallet_transactions
   has_many :recurring_transaction_rules
 
-  monetize :balance_cents, :ongoing_balance_cents, :ongoing_usage_balance_cents
+  monetize :balance_cents
   monetize :consumed_amount_cents
+  monetize :ongoing_balance_cents, :ongoing_usage_balance_cents, with_model_currency: :balance_currency
 
   validates :rate_amount, numericality: {greater_than: 0}
 

app/services/events/post_process_service.rb

@@ -25,7 +25,7 @@ def call
 
       expire_cached_charges(subscriptions)
       flag_refresh_from_subscription
-      flag_wallets_for_refresh
+      customer&.flag_wallets_for_refresh
 
       handle_pay_in_advance
 
@@ -99,13 +99,6 @@ def flag_refresh_from_subscription
       subscriptions.select(&:active?).each { |s| LifetimeUsages::FlagRefreshFromSubscriptionService.new(subscription: s).call }
     end
 
-    def flag_wallets_for_refresh
-      return unless customer
-      return unless customer.wallets.active.any?
-
-      customer.wallets.active.update_all(ready_to_be_refreshed: true) # rubocop:disable Rails/SkipsModelValidations
-    end
-
     def handle_pay_in_advance
       return unless billable_metric
       return unless charges.any?

app/services/invoices/refresh_draft_service.rb

@@ -31,6 +31,7 @@ def call
 
       ActiveRecord::Base.transaction do
         invoice.update!(ready_to_be_refreshed: false) if invoice.ready_to_be_refreshed?
+        old_total_amount_cents = invoice.total_amount_cents
 
         old_fee_values = invoice_credit_note_items.map do |item|
           {credit_note_item_id: item.id, fee_amount_cents: item.fee&.amount_cents}
@@ -67,7 +68,10 @@ def call
         end
         calculate_result.raise_if_error!
 
-        flag_lifetime_usage_for_refresh
+        if old_total_amount_cents != invoice.total_amount_cents
+          flag_lifetime_usage_for_refresh
+          invoice.customer.flag_wallets_for_refresh
+        end
 
         # NOTE: In case of a refresh the same day of the termination.
         invoice.fees.update_all(created_at: invoice.created_at) # rubocop:disable Rails/SkipsModelValidations

app/services/invoices/subscription_service.rb

@@ -44,6 +44,7 @@ def call
         invoice.reload
 
         flag_lifetime_usage_for_refresh
+        customer.flag_wallets_for_refresh if grace_period?
         fee_result
       end
       result.non_invoiceable_fees = fee_result.non_invoiceable_fees

app/services/wallets/balance/update_ongoing_service.rb

@@ -33,7 +33,7 @@ def call
 
       def compute_update_params
         params = {
-          ongoing_usage_balance_cents: total_usage_amount_cents,
+          ongoing_usage_balance_cents:,
           credits_ongoing_usage_balance:,
           ongoing_balance_cents:,
           credits_ongoing_balance:,
@@ -53,12 +53,16 @@ def currency
         @currency ||= wallet.ongoing_balance.currency
       end
 
+      def ongoing_usage_balance_cents
+        @ongoing_usage_balance_cents ||= total_usage_amount_cents + wallet.customer.invoices.draft.sum(:total_amount_cents)
+      end
+
       def credits_ongoing_usage_balance
-        total_usage_amount_cents.to_f.fdiv(currency.subunit_to_unit).fdiv(wallet.rate_amount)
+        ongoing_usage_balance_cents.to_f.fdiv(currency.subunit_to_unit).fdiv(wallet.rate_amount)
       end
 
       def ongoing_balance_cents
-        wallet.balance_cents - total_usage_amount_cents + pay_in_advance_usage_amount_cents
+        @ongoing_balance_cents ||= wallet.balance_cents - ongoing_usage_balance_cents + pay_in_advance_usage_amount_cents
       end
 
       def credits_ongoing_balance

spec/models/customer_spec.rb

@@ -613,4 +613,30 @@
         .and change(customer, :last_dunning_campaign_attempt_at).to(nil)
     end
   end
+
+  describe "#flag_wallets_for_refresh" do
+    context "without any wallets" do
+      it "returns nil" do
+        expect(customer.flag_wallets_for_refresh).to be_nil
+      end
+    end
+
+    context "without active wallets" do
+      it "does not flag wallets for refresh" do
+        wallet = create(:wallet, :terminated, customer:)
+
+        expect { customer.flag_wallets_for_refresh }.not_to change {
+          wallet.reload.ready_to_be_refreshed
+        }.from(false)
+      end
+    end
+
+    it "flags all active wallets for refresh" do
+      wallet = create(:wallet, customer:)
+
+      expect { customer.flag_wallets_for_refresh }.to change {
+        wallet.reload.ready_to_be_refreshed
+      }.from(false).to(true)
+    end
+  end
 end

spec/services/invoices/subscription_service_spec.rb

@@ -293,6 +293,12 @@
 
         expect(lifetime_usage.reload.recalculate_invoiced_usage).to be(false)
       end
+
+      it 'flags wallets for refresh' do
+        wallet = create(:wallet, customer:)
+
+        expect { invoice_service.call }.to change { wallet.reload.ready_to_be_refreshed }.from(false).to(true)
+      end
     end
 
     context 'when invoice already exists' do

spec/services/wallets/balance/update_ongoing_service_spec.rb

@@ -1,13 +1,16 @@
 # frozen_string_literal: true
 
-require 'rails_helper'
+require "rails_helper"
 
 RSpec.describe Wallets::Balance::UpdateOngoingService, type: :service do
   subject(:update_service) { described_class.new(wallet:, total_usage_amount_cents:, pay_in_advance_usage_amount_cents:) }
 
+  let(:organization) { create(:organization) }
+  let(:customer) { create(:customer, organization:) }
   let(:wallet) do
     create(
       :wallet,
+      customer:,
       balance_cents: 1000,
       ongoing_balance_cents: 800,
       ongoing_usage_balance_cents: 200,
@@ -22,41 +25,43 @@
 
   before { wallet }
 
-  describe '#call' do
-    it 'updates wallet balance' do
+  describe "#call" do
+    it "updates wallet balance" do
+      create(:invoice, :draft, customer:, organization:, total_amount_cents: 150)
+
       expect { update_service.call }
-        .to change(wallet.reload, :ongoing_usage_balance_cents).from(200).to(450)
-        .and change(wallet, :credits_ongoing_usage_balance).from(2.0).to(4.5)
-        .and change(wallet, :ongoing_balance_cents).from(800).to(550)
-        .and change(wallet, :credits_ongoing_balance).from(8.0).to(5.5)
+        .to change(wallet.reload, :ongoing_usage_balance_cents).from(200).to(600)
+        .and change(wallet, :credits_ongoing_usage_balance).from(2.0).to(6.0)
+        .and change(wallet, :ongoing_balance_cents).from(800).to(400)
+        .and change(wallet, :credits_ongoing_balance).from(8.0).to(4.0)
         .and change(wallet, :ready_to_be_refreshed).from(true).to(false)
 
       expect(wallet).not_to be_depleted_ongoing_balance
     end
 
-    context 'when usage amount is greater than the balance' do
+    context "when usage amount is greater than the balance" do
       let(:total_usage_amount_cents) { 1500 }
 
-      it 'updates wallet ongoing balance to a negative value' do
+      it "updates wallet ongoing balance to a negative value" do
         expect { update_service.call }
           .to change(wallet.reload, :ongoing_usage_balance_cents).from(200).to(1500)
           .and change(wallet, :credits_ongoing_usage_balance).from(2.0).to(15)
           .and change(wallet, :ongoing_balance_cents).from(800).to(-500)
           .and change(wallet, :credits_ongoing_balance).from(8.0).to(-5.0)
       end
 
-      it 'sets depleted_ongoing_balance to true' do
+      it "sets depleted_ongoing_balance to true" do
         expect { update_service.call }
           .to change(wallet.reload, :depleted_ongoing_balance).from(false).to(true)
 
         expect { update_service.call }
           .not_to change(wallet.reload, :depleted_ongoing_balance).from(true)
       end
 
-      it 'sends depleted_ongoing_balance webhook' do
+      it "sends depleted_ongoing_balance webhook" do
         expect { update_service.call }
           .to have_enqueued_job(SendWebhookJob)
-          .with('wallet.depleted_ongoing_balance', Wallet)
+          .with("wallet.depleted_ongoing_balance", Wallet)
       end
     end
   end