killbill-aplcache

diff --git a/overdue/pom.xml b/overdue/pom.xml
index 93f2878..759b56a 100644
--- a/overdue/pom.xml
+++ b/overdue/pom.xml
@@ -63,6 +63,17 @@
             <scope>test</scope>
         </dependency>
         <dependency>
+            <groupId>com.ning.billing</groupId>
+            <artifactId>killbill-catalog</artifactId>
+            <type>test-jar</type>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.ning.billing</groupId>
+            <artifactId>killbill-catalog</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
             <groupId>mysql</groupId>
             <artifactId>mysql-connector-java</artifactId>
             <scope>runtime</scope>

diff --git a/overdue/src/main/java/com/ning/billing/overdue/applicator/OverdueStateApplicatorBundle.java b/overdue/src/main/java/com/ning/billing/overdue/applicator/OverdueStateApplicatorBundle.java
index da9aa11..11cbaad 100644
--- a/overdue/src/main/java/com/ning/billing/overdue/applicator/OverdueStateApplicatorBundle.java
+++ b/overdue/src/main/java/com/ning/billing/overdue/applicator/OverdueStateApplicatorBundle.java
@@ -33,7 +33,7 @@ public class OverdueStateApplicatorBundle implements OverdueStateApplicator<Subs
     }
 
     @Override
-    public void apply(SubscriptionBundle overdueable,
+    public void apply(SubscriptionBundle bundle,
             OverdueState<SubscriptionBundle> previousOverdueState,
             OverdueState<SubscriptionBundle> nextOverdueState, DateTime timeOfNextCheck) {
         
@@ -41,41 +41,41 @@ public class OverdueStateApplicatorBundle implements OverdueStateApplicator<Subs
             return; // nothing to do
         }
         
-        cancelBundle(overdueable, previousOverdueState, nextOverdueState);
-        storeNewState(overdueable, nextOverdueState);
+        cancelBundle(bundle, previousOverdueState, nextOverdueState);
+        storeNewState(bundle, nextOverdueState);
   
         if(timeOfNextCheck != null && !nextOverdueState.isClearState()) {
-            createFutureNotification(overdueable, timeOfNextCheck);
+            createFutureNotification(bundle, timeOfNextCheck);
         }
 
         if(nextOverdueState.isClearState()) {
-            clear(overdueable);
+            clear(bundle);
         }
         
         //If new state is clear state reset next events and override table
         throw new NotImplementedException();
     }
 
-    private void cancelBundle(SubscriptionBundle overdueable,
+    private void cancelBundle(SubscriptionBundle bundle,
             OverdueState<SubscriptionBundle> previousOverdueState,
             OverdueState<SubscriptionBundle> nextOverdueState) {
         // TODO Auto-generated method stub
         
     }
 
-    private void storeNewState(SubscriptionBundle overdueable,
+    private void storeNewState(SubscriptionBundle bundle,
             OverdueState<SubscriptionBundle> nextOverdueState) {
         // TODO Auto-generated method stub
         
     }
 
-    private void createFutureNotification(SubscriptionBundle overdueable,
+    private void createFutureNotification(SubscriptionBundle bundle,
             DateTime timeOfNextCheck) {
         // TODO Auto-generated method stub
         
     }
 
-    private void clear(SubscriptionBundle overdueable) {
+    private void clear(SubscriptionBundle bundle) {
         // TODO Clear any future events plus overrides
         
     }

diff --git a/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculator.java b/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculator.java
index 1ec10a8..a5d66aa 100644
--- a/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculator.java
+++ b/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculator.java
@@ -37,6 +37,18 @@ public abstract class BillingStateCalculator<T extends Overdueable> {
 
     private final InvoiceUserApi invoiceApi;
     private final Clock clock;
+    
+    protected class InvoiceDateComparator implements Comparator<Invoice> {
+        @Override
+        public int compare(Invoice i1, Invoice i2) {
+            DateTime d1 = i1.getInvoiceDate();
+            DateTime d2 = i2.getInvoiceDate();
+            if(d1.compareTo(d2) == 0) {
+                return i1.hashCode() - i2.hashCode(); // consistent (arbitrary) resolution for tied dates
+            }
+            return d1.compareTo(d2);
+        }
+    }
 
     @Inject 
     public BillingStateCalculator(InvoiceUserApi invoiceApi, Clock clock) {
@@ -59,20 +71,9 @@ public abstract class BillingStateCalculator<T extends Overdueable> {
         return sum;
     }
 
-    protected SortedSet<Invoice> unpaidInvoicesFor(UUID id) {
-        Collection<Invoice> invoices = invoiceApi.getUnpaidInvoicesByAccountId(id, clock.getUTCNow());
-        SortedSet<Invoice> sortedInvoices = new TreeSet<Invoice>(new Comparator<Invoice>() {
-            @Override
-            public int compare(Invoice i1, Invoice i2) {
-                DateTime d1 = i1.getInvoiceDate();
-                DateTime d2 = i2.getInvoiceDate();
-                if(d1.compareTo(d2) == 0) {
-                    return i1.hashCode() - i2.hashCode(); // consistent (arbitrary) resolution for tied dates
-                }
-                return d1.compareTo(d2);
-            }
-            
-        });
+    protected SortedSet<Invoice> unpaidInvoicesForAccount(UUID accountId) {
+        Collection<Invoice> invoices = invoiceApi.getUnpaidInvoicesByAccountId(accountId, clock.getUTCNow());
+        SortedSet<Invoice> sortedInvoices = new TreeSet<Invoice>(new InvoiceDateComparator());
         sortedInvoices.addAll(invoices);
         return sortedInvoices;
     }

diff --git a/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculatorBundle.java b/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculatorBundle.java
index d7bbdd7..cbd9219 100644
--- a/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculatorBundle.java
+++ b/overdue/src/main/java/com/ning/billing/overdue/calculator/BillingStateCalculatorBundle.java
@@ -18,6 +18,7 @@ package com.ning.billing.overdue.calculator;
 
 import java.math.BigDecimal;
 import java.util.SortedSet;
+import java.util.TreeSet;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
@@ -27,13 +28,13 @@ import com.ning.billing.catalog.api.BillingPeriod;
 import com.ning.billing.catalog.api.PhaseType;
 import com.ning.billing.catalog.api.PriceList;
 import com.ning.billing.catalog.api.Product;
-import com.ning.billing.catalog.api.overdue.BillingState;
 import com.ning.billing.catalog.api.overdue.BillingStateBundle;
 import com.ning.billing.catalog.api.overdue.PaymentResponse;
 import com.ning.billing.entitlement.api.overdue.EntitlementOverdueApi;
 import com.ning.billing.entitlement.api.user.Subscription;
 import com.ning.billing.entitlement.api.user.SubscriptionBundle;
 import com.ning.billing.invoice.api.Invoice;
+import com.ning.billing.invoice.api.InvoiceItem;
 import com.ning.billing.invoice.api.InvoiceUserApi;
 import com.ning.billing.util.clock.Clock;
 import com.ning.billing.util.tag.Tag;
@@ -49,9 +50,9 @@ public class BillingStateCalculatorBundle  extends BillingStateCalculator<Subscr
     }
     
     @Override
-    public BillingState<SubscriptionBundle> calculateBillingState(SubscriptionBundle bundle) {
+    public BillingStateBundle calculateBillingState(SubscriptionBundle bundle) {
         
-        SortedSet<Invoice> unpaidInvoices = unpaidInvoicesFor(bundle.getId());
+        SortedSet<Invoice> unpaidInvoices = unpaidInvoicesForBundle(bundle.getId(), bundle.getAccountId());
  
         Subscription basePlan = entitlementApi.getBaseSubscription(bundle.getId());
         
@@ -80,4 +81,27 @@ public class BillingStateCalculatorBundle  extends BillingStateCalculator<Subscr
             basePlanPhaseType);
         
     }
+
+    public SortedSet<Invoice> unpaidInvoicesForBundle(UUID bundleId, UUID accountId) {
+        SortedSet<Invoice> unpaidInvoices = unpaidInvoicesForAccount(accountId);
+        SortedSet<Invoice> result = new TreeSet<Invoice>(new InvoiceDateComparator());
+        result.addAll(unpaidInvoices);
+        for(Invoice invoice : unpaidInvoices) {
+            if(!invoiceHasAnItemFromBundle(invoice, bundleId)) {
+                result.remove(invoice);
+            }
+        }
+        return result;
+    }
+
+    private boolean invoiceHasAnItemFromBundle(Invoice invoice, UUID bundleId) {
+        for(InvoiceItem item : invoice.getInvoiceItems()) {
+            if(item.getBundleId().equals(bundleId)) {
+                return true;
+            }
+        }
+        return false;
+    }
+    
+    
 }

diff --git a/overdue/src/test/java/com/ning/billing/overdue/calculator/TestBillingStateCalculator.java b/overdue/src/test/java/com/ning/billing/overdue/calculator/TestBillingStateCalculator.java
index 0ccbd7d..8d2a525 100644
--- a/overdue/src/test/java/com/ning/billing/overdue/calculator/TestBillingStateCalculator.java
+++ b/overdue/src/test/java/com/ning/billing/overdue/calculator/TestBillingStateCalculator.java
@@ -19,6 +19,7 @@ package com.ning.billing.overdue.calculator;
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.List;
 import java.util.SortedSet;
 import java.util.UUID;
 
@@ -29,6 +30,7 @@ import org.testng.annotations.Test;
 import com.ning.billing.catalog.api.overdue.BillingState;
 import com.ning.billing.entitlement.api.user.SubscriptionBundle;
 import com.ning.billing.invoice.api.Invoice;
+import com.ning.billing.invoice.api.InvoiceItem;
 import com.ning.billing.invoice.api.InvoiceUserApi;
 import com.ning.billing.mock.BrainDeadProxyFactory;
 import com.ning.billing.mock.BrainDeadProxyFactory.ZombieControl;
@@ -44,9 +46,9 @@ public class TestBillingStateCalculator {
     public BillingStateCalculator<SubscriptionBundle> createBSCalc() {
         now = new DateTime();
         Collection<Invoice> invoices = new ArrayList<Invoice>();
-        invoices.add(createInvoice(now, BigDecimal.ZERO));
-        invoices.add(createInvoice(now.plusDays(1), BigDecimal.TEN));
-        invoices.add(createInvoice(now.plusDays(2), new BigDecimal("100.0")));
+        invoices.add(createInvoice(now, BigDecimal.ZERO, null));
+        invoices.add(createInvoice(now.plusDays(1), BigDecimal.TEN, null));
+        invoices.add(createInvoice(now.plusDays(2), new BigDecimal("100.0"), null));
      
         ((ZombieControl)invoiceApi).addResult("getUnpaidInvoicesByAccountId", invoices);
             
@@ -58,11 +60,12 @@ public class TestBillingStateCalculator {
             }};
     }
     
-    public Invoice createInvoice(DateTime date, BigDecimal balance) {
+    public Invoice createInvoice(DateTime date, BigDecimal balance, List<InvoiceItem> invoiceItems) {
         Invoice invoice = BrainDeadProxyFactory.createBrainDeadProxyFor(Invoice.class);
         ((ZombieControl)invoice).addResult("getBalance", balance);
         ((ZombieControl)invoice).addResult("getInvoiceDate", date);
         ((ZombieControl)invoice).addResult("hashCode", hash++);
+        ((ZombieControl)invoice).addResult("getInvoiceItems", invoiceItems);
         
         return invoice;
     }
@@ -70,7 +73,7 @@ public class TestBillingStateCalculator {
     @Test(groups={"fast"}, enabled=true)
     public void testUnpaidInvoices() {
         BillingStateCalculator<SubscriptionBundle> calc = createBSCalc();
-        SortedSet<Invoice> invoices = calc.unpaidInvoicesFor(new UUID(0L,0L));
+        SortedSet<Invoice> invoices = calc.unpaidInvoicesForAccount(new UUID(0L,0L));
         
         Assert.assertEquals(invoices.size(), 3);
         Assert.assertEquals(BigDecimal.ZERO.compareTo(invoices.first().getBalance()), 0);
@@ -81,7 +84,7 @@ public class TestBillingStateCalculator {
     public void testSum() {
         
         BillingStateCalculator<SubscriptionBundle> calc = createBSCalc();
-        SortedSet<Invoice> invoices = calc.unpaidInvoicesFor(new UUID(0L,0L));
+        SortedSet<Invoice> invoices = calc.unpaidInvoicesForAccount(new UUID(0L,0L));
         Assert.assertEquals(new BigDecimal("110.0").compareTo(calc.sumBalance(invoices)), 0);
     }
 
@@ -89,7 +92,7 @@ public class TestBillingStateCalculator {
     public void testEarliest() {
         
         BillingStateCalculator<SubscriptionBundle> calc = createBSCalc();
-        SortedSet<Invoice> invoices = calc.unpaidInvoicesFor(new UUID(0L,0L));
+        SortedSet<Invoice> invoices = calc.unpaidInvoicesForAccount(new UUID(0L,0L));
         Assert.assertEquals(calc.earliest(invoices), now);
     }