killbill-memoizeit

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 0afa55c..6e65236 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
@@ -439,8 +439,13 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
                 BigDecimal accountBalance = BigDecimal.ZERO;
                 final List<InvoiceModelDao> invoices = invoiceDaoHelper.getAllInvoicesByAccountFromTransaction(invoicesTags, entitySqlDaoWrapperFactory, context);
                 for (final InvoiceModelDao cur : invoices) {
-                    accountBalance = accountBalance.add((new DefaultInvoice(cur)).getBalance());
-                    cba = cba.add(InvoiceModelDaoHelper.getCBAAmount(cur));
+
+                    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));
+                    }
                 }
                 return accountBalance.subtract(cba);
             }

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 89f60a5..76eac48 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
@@ -251,7 +251,7 @@ public class DefaultInvoice extends EntityBase implements Invoice, Cloneable {
                InvoiceCalculatorUtils.computeInvoiceBalance(currency, invoiceItems, payments, isWrittenOff() || isMigrationInvoice());
     }
 
-    private boolean hasZeroParentBalance() {
+    public boolean hasZeroParentBalance() {
         return (parentInvoice != null) && (parentInvoice.getBalance().compareTo(BigDecimal.ZERO) == 0);
     }
 

diff --git a/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestDefaultInvoiceUserApi.java b/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
index 7daf146..aa1ae0b 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/api/user/TestDefaultInvoiceUserApi.java
@@ -364,6 +364,10 @@ public class TestDefaultInvoiceUserApi extends InvoiceTestSuiteWithEmbeddedDB {
         Assert.assertEquals(creditInvoice.getStatus(), InvoiceStatus.DRAFT);
         Assert.assertEquals(creditInvoiceItem.getInvoiceId(), creditInvoice.getId());
 
+        // Verify DRAFT invoice is not taken into consideration when computing accountBalance
+        final BigDecimal accountBalance2 = invoiceUserApi.getAccountBalance(accountId, callContext);
+        Assert.assertEquals(accountBalance2, accountBalance);
+
         // move invoice from DRAFT to COMMITTED
         invoiceUserApi.commitInvoice(creditInvoice.getId(), callContext);
         creditInvoice = invoiceUserApi.getInvoice(invoiceId, callContext);