killbill-memoizeit

invoice: Rework the logic for invoice/account balance. See

5/31/2017 11:35:07 PM

Details

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/calculator/InvoiceCalculatorUtils.java b/invoice/src/main/java/org/killbill/billing/invoice/calculator/InvoiceCalculatorUtils.java
index c7af0ef..88a88c3 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/calculator/InvoiceCalculatorUtils.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/calculator/InvoiceCalculatorUtils.java
@@ -69,13 +69,10 @@ public abstract class InvoiceCalculatorUtils {
                InvoiceItemType.RECURRING.equals(invoiceItem.getInvoiceItemType());
     }
 
-    public static BigDecimal computeInvoiceBalance(final Currency currency,
-                                                   @Nullable final Iterable<InvoiceItem> invoiceItems,
-                                                   @Nullable final Iterable<InvoicePayment> invoicePayments,
-                                                   boolean writtenOffOrMigrated) {
-        if (writtenOffOrMigrated) {
-            return BigDecimal.ZERO;
-        }
+    public static BigDecimal computeRawInvoiceBalance(final Currency currency,
+                                                      @Nullable final Iterable<InvoiceItem> invoiceItems,
+                                                      @Nullable final Iterable<InvoicePayment> invoicePayments) {
+
 
         final BigDecimal amountPaid = computeInvoiceAmountPaid(currency, invoicePayments)
                 .add(computeInvoiceAmountRefunded(currency, invoicePayments));
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/CBADao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/CBADao.java
index 48de156..162c204 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/CBADao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/CBADao.java
@@ -64,7 +64,7 @@ public class CBADao {
         if (balance.compareTo(BigDecimal.ZERO) < 0) {
             // Current balance is negative, we need to generate a credit (positive CBA amount)
             return buildCBAItem(invoice, balance, context);
-        } else if (balance.compareTo(BigDecimal.ZERO) > 0 && invoice.getStatus() == InvoiceStatus.COMMITTED) {
+        } else if (balance.compareTo(BigDecimal.ZERO) > 0 && invoice.getStatus() == InvoiceStatus.COMMITTED && !invoice.isWrittenOff()) {
             // Current balance is positive and the invoice is COMMITTED, we need to use some of the existing if available (negative CBA amount)
             // PERF: in some codepaths, the CBA maybe have already been computed
             BigDecimal accountCBA = accountCBAOrNull;
@@ -86,7 +86,7 @@ public class CBADao {
     private BigDecimal getInvoiceBalance(final InvoiceModelDao invoice) {
 
         final InvoiceModelDao parentInvoice = invoice.getParentInvoice();
-        if ((parentInvoice != null) && (InvoiceModelDaoHelper.getBalance(parentInvoice).compareTo(BigDecimal.ZERO) == 0)) {
+        if ((parentInvoice != null) && (InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(parentInvoice).compareTo(BigDecimal.ZERO) == 0)) {
             final Iterable<InvoiceItemModelDao> items = Iterables.filter(parentInvoice.getInvoiceItems(), new Predicate<InvoiceItemModelDao>() {
                 @Override
                 public boolean apply(@Nullable final InvoiceItemModelDao input) {
@@ -105,7 +105,7 @@ public class CBADao {
 
         }
 
-        return InvoiceModelDaoHelper.getBalance(invoice);
+        return InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice);
     }
 
     // We let the code below rehydrate the invoice before we can add the CBA item
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
index 6e65236..bd35813 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
@@ -440,12 +440,21 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                 final List<InvoiceModelDao> invoices = invoiceDaoHelper.getAllInvoicesByAccountFromTransaction(invoicesTags, entitySqlDaoWrapperFactory, context);
                 for (final InvoiceModelDao cur : invoices) {
 
-                    final DefaultInvoice curInvoice = new DefaultInvoice(cur);
-                    final boolean ignoreForAccountBalanceComputation =  (curInvoice.getStatus().equals(InvoiceStatus.DRAFT) || curInvoice.hasZeroParentBalance());
-                    if (!ignoreForAccountBalanceComputation) {
-                        accountBalance = accountBalance.add(InvoiceModelDaoHelper.getBalance(cur));
-                        cba = cba.add(InvoiceModelDaoHelper.getCBAAmount(cur));
+                    // Skip DRAFT invoices
+                    if (cur.getStatus().equals(InvoiceStatus.DRAFT)) {
+                        continue;
                     }
+
+                    final boolean hasZeroParentBalance =
+                            cur.getParentInvoice() != null &&
+                            (cur.getParentInvoice().isWrittenOff() ||
+                             cur.getParentInvoice().getStatus() == InvoiceStatus.DRAFT ||
+                             InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(cur.getParentInvoice()).compareTo(BigDecimal.ZERO) == 0);
+
+
+                    // invoices that are WRITTEN_OFF or paid children invoices are excluded from balance computation but the cba summation needs to be included
+                    accountBalance = cur.isWrittenOff() || hasZeroParentBalance ? BigDecimal.ZERO : accountBalance.add(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(cur));
+                    cba = cba.add(InvoiceModelDaoHelper.getCBAAmount(cur));
                 }
                 return accountBalance.subtract(cba);
             }
@@ -870,7 +879,10 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
                 // Retrieve the invoice and make sure it belongs to the right account
                 final InvoiceModelDao invoice = transactional.getById(invoiceId.toString(), context);
-                if (invoice == null || !invoice.getAccountId().equals(accountId)) {
+                if (invoice == null ||
+                    !invoice.getAccountId().equals(accountId) ||
+                    invoice.isMigrated() ||
+                    invoice.getStatus() == InvoiceStatus.DRAFT) {
                     throw new InvoiceApiException(ErrorCode.INVOICE_NOT_FOUND, invoiceId);
                 }
 
@@ -889,7 +901,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
                 // Verify the final invoice balance is not negative
                 invoiceDaoHelper.populateChildren(invoice, invoicesTags, entitySqlDaoWrapperFactory, context);
-                if (InvoiceModelDaoHelper.getBalance(invoice).compareTo(BigDecimal.ZERO) < 0) {
+                if (InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).compareTo(BigDecimal.ZERO) < 0) {
                     throw new InvoiceApiException(ErrorCode.INVOICE_WOULD_BE_NEGATIVE);
                 }
 
@@ -1094,9 +1106,10 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
     private void notifyBusOfInvoiceCreation(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final InvoiceModelDao invoice, final InternalCallContext context) {
         try {
-            final BigDecimal balance = InvoiceModelDaoHelper.getBalance(invoice);
+            // This is called for a new COMMITTED invoice (which cannot be writtenOff as it does not exist yet, so rawBalance == balance)
+            final BigDecimal rawBalance = InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice);
             final DefaultInvoiceCreationEvent event = new DefaultInvoiceCreationEvent(invoice.getId(), invoice.getAccountId(),
-                                                                                                            balance, invoice.getCurrency(),
+                                                                                      rawBalance, invoice.getCurrency(),
                                                                                                             context.getAccountRecordId(), context.getTenantRecordId(),
                                                                                                             context.getUserToken());
             eventBus.postFromTransaction(event, entitySqlDaoWrapperFactory.getHandle().getConnection());
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDaoHelper.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDaoHelper.java
index cfd92a1..ecebd92 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDaoHelper.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDaoHelper.java
@@ -175,9 +175,10 @@ public class InvoiceDaoHelper {
             @Override
             public boolean apply(final InvoiceModelDao in) {
                 final InvoiceModelDao invoice = (in.getParentInvoice() == null) ? in : in.getParentInvoice();
-                final BigDecimal balance = InvoiceModelDaoHelper.getBalance(invoice);
+                final BigDecimal balance = InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice);
                 log.debug("Computed balance={} for invoice={}", balance, in);
-                return InvoiceStatus.COMMITTED.equals(in.getStatus()) && (balance.compareTo(BigDecimal.ZERO) >= 1) &&
+                return InvoiceStatus.COMMITTED.equals(in.getStatus()) &&
+                       (balance.compareTo(BigDecimal.ZERO) >= 1 && !in.isWrittenOff()) &&
                        (upToDate == null || in.getTargetDate() == null || !in.getTargetDate().isAfter(upToDate));
             }
         });
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceModelDaoHelper.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceModelDaoHelper.java
index a240128..877e741 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceModelDaoHelper.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceModelDaoHelper.java
@@ -18,8 +18,6 @@ package org.killbill.billing.invoice.dao;
 
 import java.math.BigDecimal;
 
-import javax.annotation.Nullable;
-
 import org.killbill.billing.invoice.api.InvoiceItem;
 import org.killbill.billing.invoice.api.InvoicePayment;
 import org.killbill.billing.invoice.calculator.InvoiceCalculatorUtils;
@@ -33,22 +31,25 @@ public class InvoiceModelDaoHelper {
 
     private InvoiceModelDaoHelper() {}
 
-    public static BigDecimal getBalance(final InvoiceModelDao invoiceModelDao) {
-        return InvoiceCalculatorUtils.computeInvoiceBalance(invoiceModelDao.getCurrency(),
-                                                            Iterables.transform(invoiceModelDao.getInvoiceItems(), new Function<InvoiceItemModelDao, InvoiceItem>() {
-                                                                @Override
-                                                                public InvoiceItem apply(final InvoiceItemModelDao input) {
-                                                                    return InvoiceItemFactory.fromModelDao(input);
-                                                                }
-                                                            }),
-                                                            Iterables.transform(invoiceModelDao.getInvoicePayments(), new Function<InvoicePaymentModelDao, InvoicePayment>() {
-                                                                @Nullable
-                                                                @Override
-                                                                public InvoicePayment apply(final InvoicePaymentModelDao input) {
-                                                                    return new DefaultInvoicePayment(input);
-                                                                }
-                                                            }),
-                                                            invoiceModelDao.isMigrated() || invoiceModelDao.isWrittenOff());
+    public static BigDecimal getRawBalanceForRegularInvoice(final InvoiceModelDao invoiceModelDao) {
+
+        if (invoiceModelDao.isMigrated()) {
+            return BigDecimal.ZERO;
+        }
+
+        return InvoiceCalculatorUtils.computeRawInvoiceBalance(invoiceModelDao.getCurrency(),
+                                                               Iterables.transform(invoiceModelDao.getInvoiceItems(), new Function<InvoiceItemModelDao, InvoiceItem>() {
+                                                                   @Override
+                                                                   public InvoiceItem apply(final InvoiceItemModelDao input) {
+                                                                       return InvoiceItemFactory.fromModelDao(input);
+                                                                   }
+                                                               }),
+                                                               Iterables.transform(invoiceModelDao.getInvoicePayments(), new Function<InvoicePaymentModelDao, InvoicePayment>() {
+                                                                   @Override
+                                                                   public InvoicePayment apply(final InvoicePaymentModelDao input) {
+                                                                       return new DefaultInvoicePayment(input);
+                                                                   }
+                                                               }));
     }
 
     public static BigDecimal getCBAAmount(final InvoiceModelDao invoiceModelDao) {
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
index a300bf1..d8fc889 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
@@ -861,7 +861,7 @@ public class InvoiceDispatcher {
 
         if (parentInvoiceModelDao == null) {
             throw new InvoiceApiException(ErrorCode.INVOICE_MISSING_PARENT_INVOICE, childInvoiceModelDao.getId());
-        } else if (InvoiceModelDaoHelper.getBalance(parentInvoiceModelDao).compareTo(BigDecimal.ZERO) == 0) {
+        } else if (InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(parentInvoiceModelDao).compareTo(BigDecimal.ZERO) == 0) {
             // ignore item adjustments for paid invoices.
             return;
         }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/model/DefaultInvoice.java b/invoice/src/main/java/org/killbill/billing/invoice/model/DefaultInvoice.java
index 76eac48..9230643 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/model/DefaultInvoice.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/model/DefaultInvoice.java
@@ -246,9 +246,14 @@ public class DefaultInvoice extends EntityBase implements Invoice, Cloneable {
 
     @Override
     public BigDecimal getBalance() {
-        return getStatus().equals(InvoiceStatus.DRAFT) || hasZeroParentBalance() ?
-               BigDecimal.ZERO :
-               InvoiceCalculatorUtils.computeInvoiceBalance(currency, invoiceItems, payments, isWrittenOff() || isMigrationInvoice());
+        if (isWrittenOff() ||
+            isMigrationInvoice() ||
+            getStatus() == InvoiceStatus.DRAFT ||
+            hasZeroParentBalance()) {
+            return BigDecimal.ZERO;
+        } else {
+            return InvoiceCalculatorUtils.computeRawInvoiceBalance(currency, invoiceItems, payments);
+        }
     }
 
     public boolean hasZeroParentBalance() {
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java b/invoice/src/test/java/org/killbill/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
index f5849cd..4447c42 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/api/migration/TestDefaultInvoiceMigrationApi.java
@@ -81,7 +81,7 @@ public class TestDefaultInvoiceMigrationApi extends InvoiceTestSuiteWithEmbedded
         Assert.assertEquals(invoice.getInvoiceItems().size(), 1);
         Assert.assertEquals(invoice.getInvoiceItems().get(0).getAmount().compareTo(MIGRATION_INVOICE_AMOUNT), 0);
         Assert.assertEquals(invoice.getInvoiceItems().get(0).getType(), InvoiceItemType.RECURRING);
-        Assert.assertEquals(InvoiceModelDaoHelper.getBalance(invoice).compareTo(BigDecimal.ZERO), 0);
+        Assert.assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).compareTo(BigDecimal.ZERO), 0);
         Assert.assertEquals(invoice.getCurrency(), MIGRATION_INVOICE_CURRENCY);
         Assert.assertTrue(invoice.isMigrated());
 
@@ -108,7 +108,7 @@ public class TestDefaultInvoiceMigrationApi extends InvoiceTestSuiteWithEmbedded
     public void testBalance() throws InvoiceApiException {
         final InvoiceModelDao migrationInvoice = invoiceDao.getById(migrationInvoiceId, internalCallContext);
         final InvoiceModelDao regularInvoice = invoiceDao.getById(regularInvoiceId, internalCallContext);
-        final BigDecimal balanceOfAllInvoices = InvoiceModelDaoHelper.getBalance(migrationInvoice).add(InvoiceModelDaoHelper.getBalance(regularInvoice));
+        final BigDecimal balanceOfAllInvoices = InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(migrationInvoice).add(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(regularInvoice));
 
         final BigDecimal accountBalance = invoiceUserApi.getAccountBalance(accountId, callContext);
         Assert.assertEquals(accountBalance.compareTo(balanceOfAllInvoices), 0);
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestInvoiceFlagBehaviors.java b/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestInvoiceFlagBehaviors.java
new file mode 100644
index 0000000..6a8795f
--- /dev/null
+++ b/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestInvoiceFlagBehaviors.java
@@ -0,0 +1,225 @@
+/*
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.invoice.api.user;
+
+import java.math.BigDecimal;
+import java.util.List;
+import java.util.UUID;
+
+import org.killbill.billing.ObjectType;
+import org.killbill.billing.account.api.Account;
+import org.killbill.billing.invoice.InvoiceTestSuiteWithEmbeddedDB;
+import org.killbill.billing.invoice.api.Invoice;
+import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.invoice.api.InvoiceStatus;
+import org.killbill.billing.invoice.model.ExternalChargeInvoiceItem;
+import org.killbill.billing.invoice.model.FixedPriceInvoiceItem;
+import org.killbill.billing.util.tag.ControlTagType;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.google.common.collect.ImmutableList;
+
+import static org.testng.Assert.assertEquals;
+
+public class TestInvoiceFlagBehaviors extends InvoiceTestSuiteWithEmbeddedDB {
+
+    private UUID accountId;
+
+    @Override
+    @BeforeMethod(groups = "slow")
+    public void beforeMethod() throws Exception {
+        super.beforeMethod();
+        final Account account = invoiceUtil.createAccount(callContext);
+        accountId = account.getId();
+    }
+
+    @Test(groups = "slow", description = "Verify invoice/account balance with a WRITTEN_OFF invoice. Verify account credit is not used against such invoice")
+    public void testWrittenOffInvoiceBeforeAccountCredit() throws Exception {
+
+        // Create new invoice with one charge and expect account credit to be used
+        final List<InvoiceItem> items = invoiceUserApi.insertExternalCharges(accountId, null, ImmutableList.<InvoiceItem>of(new ExternalChargeInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), null, accountId, null, null, null, BigDecimal.TEN, accountCurrency)), true, callContext);
+        assertEquals(items.size(), 1);
+
+        // Check both invoice and account balance is 10.00
+        final UUID invoiceId = items.get(0).getInvoiceId();
+
+        final Invoice invoice0 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice0.getBalance().compareTo(BigDecimal.TEN), 0);
+        final BigDecimal accountBalance0 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance0.compareTo(BigDecimal.TEN), 0);
+
+        final BigDecimal accountCBA0 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA0.compareTo(BigDecimal.ZERO), 0);
+
+        // Tag invoice with WRITTEN_OFF and expect balance to now show as Zero
+        tagUserApi.addTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), callContext);
+
+        // Check both invoice and account balance is NOW  0
+        final Invoice invoice1 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice1.getBalance().compareTo(BigDecimal.ZERO), 0);
+
+        final BigDecimal accountBalance1 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance1.compareTo(BigDecimal.ZERO), 0);
+
+        final BigDecimal accountCBA1 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA1.compareTo(BigDecimal.ZERO), 0);
+
+        // Add credit on the account
+        invoiceUserApi.insertCredit(accountId, BigDecimal.TEN, null, accountCurrency, true, null, callContext);
+
+        final Invoice invoice2 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice2.getBalance().compareTo(BigDecimal.ZERO), 0);
+        assertEquals(invoice2.getInvoiceItems().size(), 1);
+
+        final BigDecimal accountBalance2 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance2.compareTo(new BigDecimal("-10.00")), 0);
+
+        final BigDecimal accountCBA2 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA2.compareTo(BigDecimal.TEN), 0);
+
+
+    }
+
+    @Test(groups = "slow", description = "Verify invoice/account balance with a WRITTEN_OFF invoice. Verify behavior when WRITTEN_OFF tag is added after credit was added to invoice" )
+    public void testWrittenOffInvoiceWithAccountCredit() throws Exception {
+
+        // Add credit on the account
+        invoiceUserApi.insertCredit(accountId, BigDecimal.TEN, null, accountCurrency, true, null, callContext);
+
+        final BigDecimal accountBalance0 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance0.compareTo(new BigDecimal("-10.0")), 0);
+
+        final BigDecimal accountCBA0 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA0.compareTo(BigDecimal.TEN), 0);
+
+        // Create new invoice with one charge and expect account credit to be used
+        final List<InvoiceItem> items = invoiceUserApi.insertExternalCharges(accountId, null, ImmutableList.<InvoiceItem>of(new ExternalChargeInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), null, accountId, null, null, null, new BigDecimal("13.5"), accountCurrency)), true, callContext);
+        assertEquals(items.size(), 1);
+
+        final BigDecimal accountBalance1 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance1.compareTo(new BigDecimal("3.5")), 0);
+
+        final BigDecimal accountCBA1 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA1.compareTo(BigDecimal.ZERO), 0);
+
+        // Tag invoice with WRITTEN_OFF and expect balance to now show as Zero
+        tagUserApi.addTag(items.get(0).getInvoiceId(), ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), callContext);
+
+        final BigDecimal accountBalance2 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance2.compareTo(BigDecimal.ZERO), 0);
+
+        final BigDecimal accountCBA2 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA2.compareTo(BigDecimal.ZERO), 0);
+    }
+
+    @Test(groups = "slow", description = "Verify invoice/account balance with a WRITTEN_OFF invoice. Verify account credit is not added to a previously WRITTEN_OFF invoice." )
+    public void testWrittenOffInvoiceWithAccountCredit2() throws Exception {
+
+        // Add credit on the account
+        invoiceUserApi.insertCredit(accountId, BigDecimal.TEN, null, accountCurrency, true, null, callContext);
+
+        final BigDecimal accountBalance0 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance0.compareTo(new BigDecimal("-10.0")), 0);
+
+        final BigDecimal accountCBA0 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA0.compareTo(BigDecimal.TEN), 0);
+
+        // Create new invoice with one charge and expect account credit to be used
+        final List<InvoiceItem> items = invoiceUserApi.insertExternalCharges(accountId, null, ImmutableList.<InvoiceItem>of(new ExternalChargeInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), null, accountId, null, null, null, new BigDecimal("4.0"), accountCurrency)), true, callContext);
+        assertEquals(items.size(), 1);
+
+        // Tag invoice with WRITTEN_OFF
+        final UUID invoiceId = items.get(0).getInvoiceId();
+        tagUserApi.addTag(invoiceId, ObjectType.INVOICE, ControlTagType.WRITTEN_OFF.getId(), callContext);
+
+        // Add another charge on the **same invoice** => Because it is WRITTEN_OFF, we expect the CBA logic to not apply any credit
+        final List<InvoiceItem> items2 = invoiceUserApi.insertExternalCharges(accountId, null, ImmutableList.<InvoiceItem>of(new ExternalChargeInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), invoiceId, accountId, null, null, null, new BigDecimal("10.0"), accountCurrency)), true, callContext);
+        assertEquals(items2.size(), 1);
+
+        final BigDecimal accountBalance2 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance2.compareTo(new BigDecimal("-6.00")), 0);
+
+        final BigDecimal accountCBA2 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA2.compareTo(new BigDecimal("6.00")), 0);
+    }
+
+
+    @Test(groups = "slow", description = "Verify invoice/account balance with migrated invoice. Verify account credit is not consumed and that invoice/account balance does not take into account migrated invoice.")
+    public void testMigratedInvoiceWithAccountCredit() throws Exception {
+
+        // Add credit on the account
+        invoiceUserApi.insertCredit(accountId, BigDecimal.TEN, null, accountCurrency, true, null, callContext);
+
+        final UUID invoiceId = invoiceUserApi.createMigrationInvoice(accountId, null, ImmutableList.<InvoiceItem>of(new FixedPriceInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), null, accountId, null, null, "foo", "bar", null, null, BigDecimal.ONE, accountCurrency)), callContext);
+
+        final Invoice invoice1 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice1.getBalance().compareTo(BigDecimal.ZERO), 0);
+
+        // Verify credit is **not applied** against migration invoice
+        final BigDecimal accountBalance0 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance0.compareTo(new BigDecimal("-10.0")), 0);
+
+        final BigDecimal accountCBA0 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA0.compareTo(BigDecimal.TEN), 0);
+
+    }
+
+    @Test(groups = "slow", description = "Verify invoice/account balance with DRAFT invoice. Verify that invoice/account balance are ZERO in DRAFT mode but becomes visible after it hasa been COMMITTED." )
+    public void testDraftInvoiceWithAccountCredit() throws Exception {
+
+        // Add credit on the account
+        invoiceUserApi.insertCredit(accountId, BigDecimal.TEN, null, accountCurrency, true, null, callContext);
+
+        // Create new invoice with one charge and expect account credit to be used
+        final List<InvoiceItem> items = invoiceUserApi.insertExternalCharges(accountId, null, ImmutableList.<InvoiceItem>of(new ExternalChargeInvoiceItem(UUID.randomUUID(), clock.getUTCNow(), null, accountId, null, null, null, new BigDecimal("4.0"), accountCurrency)), false, callContext);
+        assertEquals(items.size(), 1);
+
+        final UUID invoiceId = items.get(0).getInvoiceId();
+
+        final Invoice invoice1 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice1.getStatus(), InvoiceStatus.DRAFT);
+
+        // Verify CBA was *NOT* applied against DRAFT invoice
+        assertEquals(invoice1.getInvoiceItems().size(), 1);
+        // And balance is ZERO because DRAFT mode
+        assertEquals(invoice1.getBalance().compareTo(BigDecimal.ZERO), 0);
+
+        // Verify credit is not applied against migration invoice
+        final BigDecimal accountBalance0 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance0.compareTo(new BigDecimal("-10.0")), 0);
+
+        final BigDecimal accountCBA0 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA0.compareTo(BigDecimal.TEN), 0);
+
+        invoiceUserApi.commitInvoice(invoiceId, callContext);
+
+        final Invoice invoice2 = invoiceUserApi.getInvoice(invoiceId, callContext);
+        assertEquals(invoice2.getStatus(), InvoiceStatus.COMMITTED);
+
+        // Verify this time credit was applied against COMMITTED invoice
+        assertEquals(invoice2.getBalance().compareTo(BigDecimal.ZERO), 0);
+
+        final BigDecimal accountBalance1 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        assertEquals(accountBalance1.compareTo(new BigDecimal("-6.0")), 0);
+
+        final BigDecimal accountCBA1 = invoiceUserApi.getAccountCBA(accountId, callContext);
+        assertEquals(accountCBA1.compareTo(new BigDecimal("6.0")), 0);
+    }
+
+}
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java b/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
index 8ddb1f2..0e1ea27 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
@@ -68,7 +68,7 @@ public class MockInvoiceDao extends MockEntityDaoBase<InvoiceModelDao, Invoice, 
         }
         try {
             eventBus.post(new DefaultInvoiceCreationEvent(invoice.getId(), invoice.getAccountId(),
-                                                          InvoiceModelDaoHelper.getBalance(invoice), invoice.getCurrency(),
+                                                          InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice), invoice.getCurrency(),
                                                           context.getAccountRecordId(), context.getTenantRecordId(), context.getUserToken()));
         } catch (final PersistentBus.EventBusException ex) {
             throw new RuntimeException(ex);
@@ -260,7 +260,7 @@ public class MockInvoiceDao extends MockEntityDaoBase<InvoiceModelDao, Invoice, 
 
         for (final InvoiceModelDao invoice : getAll(context)) {
             if (accountId.equals(invoice.getAccountId())) {
-                balance = balance.add(InvoiceModelDaoHelper.getBalance(invoice));
+                balance = balance.add(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice));
             }
         }
 
@@ -272,7 +272,7 @@ public class MockInvoiceDao extends MockEntityDaoBase<InvoiceModelDao, Invoice, 
         final List<InvoiceModelDao> unpaidInvoices = new ArrayList<InvoiceModelDao>();
 
         for (final InvoiceModelDao invoice : getAll(context)) {
-            if (accountId.equals(invoice.getAccountId()) && (InvoiceModelDaoHelper.getBalance(invoice).compareTo(BigDecimal.ZERO) > 0) && !invoice.isMigrated()) {
+            if (accountId.equals(invoice.getAccountId()) && (InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).compareTo(BigDecimal.ZERO) > 0) && !invoice.isMigrated()) {
                 unpaidInvoices.add(invoice);
             }
         }
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/dao/TestInvoiceDao.java b/invoice/src/test/java/org/killbill/billing/invoice/dao/TestInvoiceDao.java
index 3b245a3..9d7199a 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/dao/TestInvoiceDao.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/dao/TestInvoiceDao.java
@@ -127,7 +127,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
         assertTrue(thisInvoice.getInvoiceDate().compareTo(invoiceDate) == 0);
         assertEquals(thisInvoice.getCurrency(), Currency.USD);
         assertEquals(thisInvoice.getInvoiceItems().size(), 0);
-        assertTrue(InvoiceModelDaoHelper.getBalance(thisInvoice).compareTo(BigDecimal.ZERO) == 0);
+        assertTrue(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(thisInvoice).compareTo(BigDecimal.ZERO) == 0);
     }
 
     @Test(groups = "slow")
@@ -147,7 +147,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
 
         final InvoiceModelDao savedInvoice = invoiceDao.getById(invoiceId, context);
         assertNotNull(savedInvoice);
-        assertEquals(InvoiceModelDaoHelper.getBalance(savedInvoice).compareTo(new BigDecimal("21.00")), 0);
+        assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(savedInvoice).compareTo(new BigDecimal("21.00")), 0);
         assertEquals(savedInvoice.getInvoiceItems().size(), 1);
 
         final BigDecimal paymentAmount = new BigDecimal("11.00");
@@ -159,7 +159,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
         final InvoiceModelDao retrievedInvoice = invoiceDao.getById(invoiceId, context);
         assertNotNull(retrievedInvoice);
         assertEquals(retrievedInvoice.getInvoiceItems().size(), 1);
-        assertEquals(InvoiceModelDaoHelper.getBalance(retrievedInvoice).compareTo(new BigDecimal("10.00")), 0);
+        assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(retrievedInvoice).compareTo(new BigDecimal("10.00")), 0);
     }
 
     @Test(groups = "slow")
@@ -983,7 +983,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
         assertEquals(invoices.size(), 1);
 
         final InvoiceModelDao invoice = invoices.get(0);
-        assertTrue(InvoiceModelDaoHelper.getBalance(invoice).compareTo(BigDecimal.ZERO) == 0);
+        assertTrue(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).compareTo(BigDecimal.ZERO) == 0);
         final List<InvoiceItemModelDao> invoiceItems = invoice.getInvoiceItems();
         assertEquals(invoiceItems.size(), 2);
         boolean foundCredit = false;
@@ -1058,7 +1058,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
         assertEquals(invoices.size(), 1);
 
         final InvoiceModelDao invoice = invoices.get(0);
-        assertTrue(InvoiceModelDaoHelper.getBalance(invoice).compareTo(expectedBalance) == 0);
+        assertTrue(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).compareTo(expectedBalance) == 0);
         final List<InvoiceItemModelDao> invoiceItems = invoice.getInvoiceItems();
         assertEquals(invoiceItems.size(), expectCBA ? 3 : 2);
         boolean foundCredit = false;
@@ -1236,10 +1236,10 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
         invoiceUtil.createInvoice(invoice2, context);
 
         final InvoiceModelDao savedInvoice1 = invoiceDao.getById(invoice1.getId(), context);
-        assertEquals(InvoiceModelDaoHelper.getBalance(savedInvoice1), KillBillMoney.of(TEN, savedInvoice1.getCurrency()));
+        assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(savedInvoice1), KillBillMoney.of(TEN, savedInvoice1.getCurrency()));
 
         final InvoiceModelDao savedInvoice2 = invoiceDao.getById(invoice2.getId(), context);
-        assertEquals(InvoiceModelDaoHelper.getBalance(savedInvoice2), KillBillMoney.of(FIVE, savedInvoice2.getCurrency()));
+        assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(savedInvoice2), KillBillMoney.of(FIVE, savedInvoice2.getCurrency()));
     }
 
     @Test(groups = "slow")
@@ -1388,7 +1388,7 @@ public class TestInvoiceDao extends InvoiceTestSuiteWithEmbeddedDB {
 
         assertNotNull(savedInvoice);
         assertEquals(savedInvoice.getInvoiceItems().size(), 2);
-        assertEquals(InvoiceModelDaoHelper.getBalance(savedInvoice).compareTo(cheapAmount), 0);
+        assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(savedInvoice).compareTo(cheapAmount), 0);
     }
 
     @Test(groups = "slow")
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
index dc90d85..e1df2df 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
@@ -316,7 +316,7 @@ public class TestInvoiceHelper {
 
     public void verifyInvoice(final UUID invoiceId, final double balance, final double cbaAmount, final InternalTenantContext context) throws InvoiceApiException {
         final InvoiceModelDao invoice = invoiceDao.getById(invoiceId, context);
-        Assert.assertEquals(InvoiceModelDaoHelper.getBalance(invoice).doubleValue(), balance);
+        Assert.assertEquals(InvoiceModelDaoHelper.getRawBalanceForRegularInvoice(invoice).doubleValue(), balance);
         Assert.assertEquals(InvoiceModelDaoHelper.getCBAAmount(invoice).doubleValue(), cbaAmount);
     }