killbill-memoizeit

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/usage/TestConsumableInArrear.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/usage/TestConsumableInArrear.java
index d77197d..fa6968b 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/usage/TestConsumableInArrear.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/usage/TestConsumableInArrear.java
@@ -186,22 +186,27 @@ public class TestConsumableInArrear extends TestIntegrationBase {
                                                 new LocalDate(2012, 5, 1),
                                                 callContext);
         assertListenerStatus();
-        final Invoice firstInvoice = invoiceChecker.checkInvoice(account.getId(), 2, callContext,
-                                                                 new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2013, 5, 1), InvoiceItemType.RECURRING, new BigDecimal("2399.95")),
-                                                                 new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, new BigDecimal("5.90")));
-        invoiceChecker.checkTrackingIds(firstInvoice, ImmutableSet.of("t1", "t2"), internalCallContext);
+        final Invoice secondInvoice = invoiceChecker.checkInvoice(account.getId(), 2, callContext,
+                                                                  new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2013, 5, 1), InvoiceItemType.RECURRING, new BigDecimal("2399.95")),
+                                                                  new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, new BigDecimal("5.90")));
+        invoiceChecker.checkTrackingIds(secondInvoice, ImmutableSet.of("t1", "t2"), internalCallContext);
 
         // Change to the Slugs plan
-        busHandler.pushExpectedEvents(NextEvent.CHANGE, NextEvent.INVOICE);
+        busHandler.pushExpectedEvents(NextEvent.CREATE, NextEvent.INVOICE);
         aoSubscription.changePlanWithDate(new DefaultEntitlementSpecifier(new PlanPhaseSpecifier("slugs-usage-in-arrear")),
                                           new LocalDate(2012, 4, 1),
                                           ImmutableList.<PluginProperty>of(),
                                           callContext);
         assertListenerStatus();
 
-        // Verify invoice
-        invoiceChecker.checkInvoice(account.getId(), 3, callContext,
-                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, BigDecimal.ZERO));
+        // Verify invoices (second invoice is unchanged)
+        final Invoice updatedSecondInvoice = invoiceChecker.checkInvoice(account.getId(), 2, callContext,
+                                                                        new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2013, 5, 1), InvoiceItemType.RECURRING, new BigDecimal("2399.95")),
+                                                                        new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, new BigDecimal("5.90")));
+        invoiceChecker.checkTrackingIds(updatedSecondInvoice, ImmutableSet.of("t1", "t2"), internalCallContext);
+        final Invoice thirdInvoice = invoiceChecker.checkInvoice(account.getId(), 3, callContext,
+                                                                  new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.USAGE, BigDecimal.ZERO));
+        invoiceChecker.checkTrackingIds(thirdInvoice, ImmutableSet.of(), internalCallContext);
 
         // Add usage data
         recordUsageData(aoSubscription.getId(), "u1", "slugs", new LocalDate(2012, 4, 1), 99L, callContext);
@@ -397,7 +402,4 @@ public class TestConsumableInArrear extends TestIntegrationBase {
 
 
     }
-
-
-
 }

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/generator/UsageInvoiceItemGenerator.java b/invoice/src/main/java/org/killbill/billing/invoice/generator/UsageInvoiceItemGenerator.java
index 5e7f5c2..a99dc62 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/generator/UsageInvoiceItemGenerator.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/generator/UsageInvoiceItemGenerator.java
@@ -196,7 +196,7 @@ public class UsageInvoiceItemGenerator extends InvoiceItemGenerator {
                     public boolean apply(final InvoiceItem input) {
                         if (input.getInvoiceItemType() == InvoiceItemType.USAGE) {
                             final Usage usage = knownUsage.get(input.getUsageName());
-                            return usage.getBillingMode() == BillingMode.IN_ARREAR;
+                            return usage != null && usage.getBillingMode() == BillingMode.IN_ARREAR;
                         }
                         return false;
                     }