killbill-memoizeit

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithInvoicePlugin.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithInvoicePlugin.java
index 17a7e68..a80c2fa 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithInvoicePlugin.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithInvoicePlugin.java
@@ -29,6 +29,7 @@ import org.awaitility.Awaitility;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
+import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.account.api.AccountData;
 import org.killbill.billing.api.TestApiListener.NextEvent;
@@ -39,6 +40,7 @@ import org.killbill.billing.catalog.api.ProductCategory;
 import org.killbill.billing.entitlement.api.DefaultEntitlement;
 import org.killbill.billing.invoice.api.DefaultInvoiceService;
 import org.killbill.billing.invoice.api.Invoice;
+import org.killbill.billing.invoice.api.InvoiceApiException;
 import org.killbill.billing.invoice.api.InvoiceItem;
 import org.killbill.billing.invoice.api.InvoiceItemType;
 import org.killbill.billing.invoice.model.ExternalChargeInvoiceItem;
@@ -61,6 +63,7 @@ import org.killbill.notificationq.api.NotificationQueueService;
 import org.killbill.notificationq.api.NotificationQueueService.NoSuchNotificationQueue;
 import org.killbill.queue.retry.RetryNotificationEvent;
 import org.killbill.queue.retry.RetryableService;
+import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -108,6 +111,11 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
     @BeforeMethod(groups = "slow")
     public void setUp() throws Exception {
         testInvoicePluginApi.additionalInvoiceItem = null;
+        testInvoicePluginApi.shouldAddTaxItem = true;
+        testInvoicePluginApi.isAborted = false;
+        testInvoicePluginApi.rescheduleDate = null;
+        testInvoicePluginApi.wasRescheduled = false;
+        testInvoicePluginApi.invocationCount = 0;
     }
 
     @Test(groups = "slow")
@@ -141,6 +149,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                                                                                    pluginLinkedItemId,
                                                                                    null);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
         // Create original subscription (Trial PHASE) -> $0 invoice but plugin added one item
         final DefaultEntitlement bpSubscription = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
         invoiceChecker.checkInvoice(account.getId(), 1, callContext,
@@ -148,6 +158,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.EXTERNAL_CHARGE, BigDecimal.TEN));
         subscriptionChecker.checkSubscriptionCreated(bpSubscription.getId(), internalCallContext);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+
         final List<Invoice> invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext);
         assertEquals(invoices.size(), 1);
         final List<InvoiceItem> invoiceItems = invoices.get(0).getInvoiceItems();
@@ -174,6 +186,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         final Account account = createAccountWithNonOsgiPaymentMethod(accountData);
         accountChecker.checkAccount(account.getId(), accountData, callContext);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
         // Create original subscription (Trial PHASE) -> $0 invoice.
         final DefaultEntitlement bpSubscription = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
         invoiceChecker.checkInvoice(account.getId(), 1, callContext,
@@ -181,6 +195,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.TAX, new BigDecimal("1.0")));
         subscriptionChecker.checkSubscriptionCreated(bpSubscription.getId(), internalCallContext);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+
         // Move to Evergreen PHASE
         busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
         clock.addDays(30);
@@ -189,6 +205,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2012, 6, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")),
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), null, InvoiceItemType.TAX, new BigDecimal("1.0")));
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+
         final List<Invoice> invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext);
         assertEquals(invoices.size(), 2);
         final InvoiceItem recurringItem = Iterables.find(invoices.get(1).getInvoiceItems(),
@@ -231,6 +249,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 6, 1), new LocalDate(2012, 7, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")),
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 6, 1), new LocalDate(2012, 6, 1), InvoiceItemType.CBA_ADJ, BigDecimal.TEN.negate()));
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 3);
+
         final List<Invoice> refreshedInvoices = invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext);
         final List<InvoiceItem> invoiceItems = refreshedInvoices.get(1).getInvoiceItems();
         final InvoiceItem invoiceItemAdjustment = Iterables.tryFind(invoiceItems, new Predicate<InvoiceItem>() {
@@ -247,6 +267,209 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
     }
 
     @Test(groups = "slow")
+    public void testAborted() throws Exception {
+        testInvoicePluginApi.shouldAddTaxItem = false;
+
+        // We take april as it has 30 days (easier to play with BCD)
+        // Set clock to the initial start date - we implicitly assume here that the account timezone is UTC
+        clock.setDay(new LocalDate(2012, 4, 1));
+
+        final AccountData accountData = getAccountData(1);
+        final Account account = createAccountWithNonOsgiPaymentMethod(accountData);
+        accountChecker.checkAccount(account.getId(), accountData, callContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
+        // Create original subscription (Trial PHASE) -> $0 invoice
+        final DefaultEntitlement bpSubscription = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        invoiceChecker.checkInvoice(account.getId(), 1, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.FIXED, new BigDecimal("0")));
+        subscriptionChecker.checkSubscriptionCreated(bpSubscription.getId(), internalCallContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+
+        // Abort invoice runs
+        testInvoicePluginApi.isAborted = true;
+
+        // Move to Evergreen PHASE
+        busHandler.pushExpectedEvents(NextEvent.PHASE);
+        clock.addDays(30);
+        assertListenerStatus();
+        assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext).size(), 1);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+
+        // No notification, so by default, the account will not be re-invoiced
+        clock.addMonths(1);
+        assertListenerStatus();
+        assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext).size(), 1);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+
+        // No notification, so by default, the account will not be re-invoiced
+        clock.addMonths(1);
+        assertListenerStatus();
+        assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext).size(), 1);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+
+        // Re-enable invoicing
+        testInvoicePluginApi.isAborted = false;
+
+        // Trigger a manual invoice run
+        busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        invoiceUserApi.triggerInvoiceGeneration(account.getId(), clock.getUTCToday(), null, callContext);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 2, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2012, 6, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")),
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 6, 1), new LocalDate(2012, 7, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")),
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 7, 1), new LocalDate(2012, 8, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 3);
+
+        // Invoicing resumes
+        busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        clock.addMonths(1);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 3, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 8, 1), new LocalDate(2012, 9, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 4);
+    }
+
+    @Test(groups = "slow")
+    public void testRescheduledViaNotification() throws Exception {
+        testInvoicePluginApi.shouldAddTaxItem = false;
+
+        // We take april as it has 30 days (easier to play with BCD)
+        // Set clock to the initial start date - we implicitly assume here that the account timezone is UTC
+        clock.setDay(new LocalDate(2012, 4, 1));
+
+        final AccountData accountData = getAccountData(1);
+        final Account account = createAccountWithNonOsgiPaymentMethod(accountData);
+        accountChecker.checkAccount(account.getId(), accountData, callContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
+        // Create original subscription (Trial PHASE) -> $0 invoice
+        final DefaultEntitlement bpSubscription = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        invoiceChecker.checkInvoice(account.getId(), 1, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.FIXED, new BigDecimal("0")));
+        subscriptionChecker.checkSubscriptionCreated(bpSubscription.getId(), internalCallContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+
+        // Reschedule invoice generation
+        final DateTime utcNow = clock.getUTCNow();
+        testInvoicePluginApi.rescheduleDate = new DateTime(2012, 5, 2, utcNow.getHourOfDay(), utcNow.getMinuteOfHour(), utcNow.getSecondOfMinute(), DateTimeZone.UTC);
+
+        // Move to Evergreen PHASE
+        busHandler.pushExpectedEvents(NextEvent.PHASE);
+        clock.addDays(30);
+        assertListenerStatus();
+        assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext).size(), 1);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+
+        // PHASE invoice has been rescheduled, reset rescheduleDate
+        testInvoicePluginApi.rescheduleDate = null;
+
+        // Move one day
+        busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        clock.addDays(1);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 2, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2012, 6, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 3);
+        Assert.assertTrue(testInvoicePluginApi.wasRescheduled);
+
+        // Invoicing resumes as expected
+        busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        clock.addDays(30);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 3, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 6, 1), new LocalDate(2012, 7, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 4);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+    }
+
+    @Test(groups = "slow")
+    public void testRescheduledViaAPI() throws Exception {
+        testInvoicePluginApi.shouldAddTaxItem = false;
+
+        // We take april as it has 30 days (easier to play with BCD)
+        // Set clock to the initial start date - we implicitly assume here that the account timezone is UTC
+        clock.setDay(new LocalDate(2012, 4, 1));
+
+        final AccountData accountData = getAccountData(1);
+        final Account account = createAccountWithNonOsgiPaymentMethod(accountData);
+        accountChecker.checkAccount(account.getId(), accountData, callContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
+        // Create original subscription (Trial PHASE) -> $0 invoice
+        final DefaultEntitlement bpSubscription = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        invoiceChecker.checkInvoice(account.getId(), 1, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.FIXED, new BigDecimal("0")));
+        subscriptionChecker.checkSubscriptionCreated(bpSubscription.getId(), internalCallContext);
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+
+        // Reschedule invoice generation at the time of the PHASE event
+        testInvoicePluginApi.rescheduleDate = new DateTime(clock.getUTCNow()).plusDays(30);
+
+        try {
+            invoiceUserApi.triggerInvoiceGeneration(account.getId(), clock.getUTCToday(), null, callContext);
+            Assert.fail();
+        } catch (final InvoiceApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.INVOICE_NOTHING_TO_DO.getCode());
+        }
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+
+        // Let the next invoice go through
+        testInvoicePluginApi.rescheduleDate = null;
+
+        // Move to Evergreen PHASE: two invoice runs will be triggers, one by SubscriptionNotificationKey (PHASE event) and one by NextBillingDateNotificationKey (reschedule)
+        busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.INVOICE, NextEvent.NULL_INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        clock.addDays(30);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 2, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 5, 1), new LocalDate(2012, 6, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 4);
+        // Cannot check wasRescheduled flag, as it would be true only for one of the runs
+
+        // Reschedule next invoice one month in the future
+        testInvoicePluginApi.rescheduleDate = clock.getUTCNow().plusMonths(1);
+        try {
+            invoiceUserApi.triggerInvoiceGeneration(account.getId(), clock.getUTCToday(), null, callContext);
+            Assert.fail();
+        } catch (final InvoiceApiException e) {
+            Assert.assertEquals(e.getCode(), ErrorCode.INVOICE_NOTHING_TO_DO.getCode());
+        }
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 5);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+
+        // Let the next invoice go through
+        testInvoicePluginApi.rescheduleDate = null;
+
+        // Move one month ahead: no NULL_INVOICE this time: since there is already a notification for that date, the reschedule is a no-op (and we keep the isRescheduled flag to false)
+        busHandler.pushExpectedEvents(NextEvent.INVOICE, NextEvent.INVOICE_PAYMENT, NextEvent.PAYMENT);
+        clock.addMonths(1);
+        assertListenerStatus();
+        invoiceChecker.checkInvoice(account.getId(), 3, callContext,
+                                    new ExpectedInvoiceItemCheck(new LocalDate(2012, 6, 1), new LocalDate(2012, 7, 1), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 6);
+        Assert.assertFalse(testInvoicePluginApi.wasRescheduled);
+    }
+
+    @Test(groups = "slow")
     public void testWithRetries() throws Exception {
         // We take april as it has 30 days (easier to play with BCD)
         // Set clock to the initial start date - we implicitly assume here that the account timezone is UTC
@@ -256,6 +479,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         final Account account = createAccountWithNonOsgiPaymentMethod(accountData);
         accountChecker.checkAccount(account.getId(), accountData, callContext);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 0);
+
         // Make invoice plugin fail
         testInvoicePluginApi.shouldThrowException = true;
 
@@ -265,6 +490,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         // Invoice failed to generate
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 0);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 1);
+
         // Verify bus event has moved to the retry service (can't easily check the timestamp unfortunately)
         // No future notification at this point (FIXED item, the PHASE event is the trigger for the next one)
         checkRetryBusEvents(1, 0);
@@ -273,6 +500,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         clock.addDeltaFromReality(5 * 60 * 1000);
         checkRetryBusEvents(2, 0);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 2);
+
         // Fix invoice plugin
         testInvoicePluginApi.shouldThrowException = false;
 
@@ -282,6 +511,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         // No notification in the main queue at this point (the PHASE event is the trigger for the next one)
         checkNotificationsNoRetry(0);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 3);
+
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 1);
         invoiceChecker.checkInvoice(account.getId(),
                                     1,
@@ -294,6 +525,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         assertListenerStatus();
         checkNotificationsNoRetry(1);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 4);
+
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 2);
         invoiceChecker.checkInvoice(account.getId(),
                                     2,
@@ -307,6 +540,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         clock.addMonths(1);
         assertListenerStatus();
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 5);
+
         // Invoice failed to generate
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 2);
 
@@ -318,6 +553,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         // Verify there are no notification duplicates
         checkRetryNotifications("2012-06-01T00:15:00", 1);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 6);
+
         // Fix invoice plugin
         testInvoicePluginApi.shouldThrowException = false;
 
@@ -326,6 +563,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         assertListenerStatus();
         checkNotificationsNoRetry(1);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 7);
+
         // Invoice was generated
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 3);
         invoiceChecker.checkInvoice(account.getId(),
@@ -340,6 +579,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         clock.setTime(new DateTime("2012-07-01T00:00:00"));
         assertListenerStatus();
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 8);
+
         // Invoice failed to generate
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false,  callContext).size(), 3);
 
@@ -353,6 +594,8 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
         assertListenerStatus();
         checkNotificationsNoRetry(1);
 
+        Assert.assertEquals(testInvoicePluginApi.invocationCount, 9);
+
         assertEquals(invoiceUserApi.getInvoicesByAccount(account.getId(), false, false, callContext).size(), 4);
     }
 
@@ -404,10 +647,28 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
 
         boolean shouldThrowException = false;
         InvoiceItem additionalInvoiceItem;
+        boolean shouldAddTaxItem = true;
+        boolean isAborted = false;
+        DateTime rescheduleDate;
+        boolean wasRescheduled = false;
+        int invocationCount = 0;
 
         @Override
         public PriorInvoiceResult priorCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> iterable) {
-            return null;
+            invocationCount++;
+            wasRescheduled = invoiceContext.isRescheduled();
+            return new PriorInvoiceResult() {
+
+                @Override
+                public boolean isAborted() {
+                    return isAborted;
+                }
+
+                @Override
+                public DateTime getRescheduleDate() {
+                    return rescheduleDate;
+                }
+            };
         }
 
         @Override
@@ -416,8 +677,10 @@ public class TestWithInvoicePlugin extends TestIntegrationBase {
                 throw new InvoicePluginApiRetryException();
             } else if (additionalInvoiceItem != null) {
                 return ImmutableList.<InvoiceItem>of(additionalInvoiceItem);
-            } else {
+            } else if (shouldAddTaxItem) {
                 return ImmutableList.<InvoiceItem>of(createTaxInvoiceItem(invoice));
+            } else {
+                return ImmutableList.<InvoiceItem>of();
             }
         }
 

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithTaxItems.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithTaxItems.java
index d0cfcee..d9b37e5 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithTaxItems.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestWithTaxItems.java
@@ -49,6 +49,7 @@ import org.killbill.billing.osgi.api.OSGIServiceDescriptor;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.util.callcontext.CallContext;
+import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -297,7 +298,7 @@ public class TestWithTaxItems extends TestIntegrationBase {
         }
 
         @Override
-        public PriorInvoiceResult priorCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> iterable) {
+        public PriorInvoiceResult priorCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> pluginProperties) {
             return null;
         }
 
@@ -307,17 +308,34 @@ public class TestWithTaxItems extends TestIntegrationBase {
             for (final TaxInvoiceItem item : taxItems) {
                 result.add(new TaxInvoiceItem(item.getId(), invoice.getId(), invoice.getAccountId(), item.getBundleId(), "Tax Item", item.getStartDate(), item.getAmount(), invoice.getCurrency()));
             }
-            taxItems.clear();
             return result;
         }
 
         @Override
-        public OnSuccessInvoiceResult onSuccessCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> iterable) {
+        public OnSuccessInvoiceResult onSuccessCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> pluginProperties) {
+            Assert.assertFalse(invoiceContext.isRescheduled());
+
+            final Invoice invoice = invoiceContext.getInvoice();
+            Assert.assertNotNull(invoice);
+            for (final TaxInvoiceItem taxInvoiceItem : taxItems) {
+                final InvoiceItem createdTaxInvoiceItem = Iterables.<InvoiceItem>find(invoice.getInvoiceItems(),
+                                                                                      new Predicate<InvoiceItem>() {
+                                                                                          @Override
+                                                                                          public boolean apply(final InvoiceItem invoiceItem) {
+                                                                                              return invoiceItem.getId().compareTo(taxInvoiceItem.getId()) == 0;
+                                                                                          }
+                                                                                      });
+                Assert.assertEquals(createdTaxInvoiceItem.getAccountId(), taxInvoiceItem.getAccountId());
+            }
+
+            reset();
+
             return null;
         }
 
         @Override
-        public OnFailureInvoiceResult onFailureCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> iterable) {
+        public OnFailureInvoiceResult onFailureCall(final InvoiceContext invoiceContext, final Iterable<PluginProperty> pluginProperties) {
+            Assert.fail();
             return null;
         }
 

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 71accaa..dcd487e 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
@@ -651,8 +651,8 @@ public class InvoiceDispatcher {
                 commitInvoiceAndSetFutureNotifications(account, null, futureAccountNotifications, internalCallContext);
             }
 
-            if (success) {
-                final DefaultInvoice refreshedInvoice = new DefaultInvoice(invoiceDao.getById(invoice.getId(), internalCallContext));
+            if (isDryRun || success) {
+                final DefaultInvoice refreshedInvoice = isDryRun ? invoice : new DefaultInvoice(invoiceDao.getById(invoice.getId(), internalCallContext));
                 invoicePluginDispatcher.onSuccessCall(targetDate, refreshedInvoice, existingInvoices, isDryRun, isRescheduled, callContext, internalCallContext);
             } else {
                 invoicePluginDispatcher.onFailureCall(targetDate, invoice, existingInvoices, isDryRun, isRescheduled, callContext, internalCallContext);

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/InvoicePluginDispatcher.java b/invoice/src/main/java/org/killbill/billing/invoice/InvoicePluginDispatcher.java
index d1abded..6bbdaf1 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoicePluginDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoicePluginDispatcher.java
@@ -25,6 +25,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.joda.time.DateTime;
@@ -102,7 +103,7 @@ public class InvoicePluginDispatcher {
     }
 
     public void onSuccessCall(final LocalDate targetDate,
-                              final DefaultInvoice invoice,
+                              @Nullable final DefaultInvoice invoice,
                               final List<Invoice> existingInvoices,
                               final boolean isDryRun,
                               final boolean isRescheduled,
@@ -113,7 +114,7 @@ public class InvoicePluginDispatcher {
     }
 
     public void onFailureCall(final LocalDate targetDate,
-                              final DefaultInvoice invoice,
+                              @Nullable final DefaultInvoice invoice,
                               final List<Invoice> existingInvoices,
                               final boolean isDryRun,
                               final boolean isRescheduled,
@@ -125,7 +126,7 @@ public class InvoicePluginDispatcher {
 
     private void onCompletionCall(final boolean isSuccess,
                                   final LocalDate targetDate,
-                                  final DefaultInvoice originalInvoice,
+                                  @Nullable final DefaultInvoice originalInvoice,
                                   final List<Invoice> existingInvoices,
                                   final boolean isDryRun,
                                   final boolean isRescheduled,
@@ -137,7 +138,7 @@ public class InvoicePluginDispatcher {
         }
 
         // We clone the original invoice so plugins don't remove/add items
-        final Invoice clonedInvoice = (Invoice) originalInvoice.clone();
+        final Invoice clonedInvoice = originalInvoice == null ? null : (Invoice) originalInvoice.clone();
         final InvoiceContext invoiceContext = new DefaultInvoiceContext(targetDate, clonedInvoice, existingInvoices, isDryRun, isRescheduled, callContext);
 
         for (final InvoicePluginApi invoicePlugin : invoicePlugins) {

diff --git a/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotificationKey.java b/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotificationKey.java
index 0e6e941..9086ccb 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotificationKey.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotificationKey.java
@@ -45,9 +45,9 @@ public class TestNextBillingDateNotificationKey {
         Assert.assertEquals(result.getUuidKey(), uuidKey);
         Assert.assertEquals(result.getTargetDate().compareTo(targetDate), 0);
         Assert.assertEquals(result.isDryRunForInvoiceNotification(), isDryRunForInvoiceNotification);
+        Assert.assertFalse(key.isRescheduled());
     }
 
-
     @Test(groups = "fast")
     public void testBasicWithUUIDKeys() throws Exception {
 
@@ -56,13 +56,14 @@ public class TestNextBillingDateNotificationKey {
         final DateTime targetDate = new DateTime();
         final Boolean isDryRunForInvoiceNotification = Boolean.FALSE;
 
-        final NextBillingDateNotificationKey key = new NextBillingDateNotificationKey(null, ImmutableList.of(uuidKey1, uuidKey2), targetDate, isDryRunForInvoiceNotification, false);
+        final NextBillingDateNotificationKey key = new NextBillingDateNotificationKey(null, ImmutableList.of(uuidKey1, uuidKey2), targetDate, isDryRunForInvoiceNotification, true);
         final String json = mapper.writeValueAsString(key);
 
         final NextBillingDateNotificationKey result = mapper.readValue(json, NextBillingDateNotificationKey.class);
         Assert.assertNull(result.getUuidKey());
         Assert.assertEquals(result.getTargetDate().compareTo(targetDate), 0);
         Assert.assertEquals(result.isDryRunForInvoiceNotification(), isDryRunForInvoiceNotification);
+        Assert.assertTrue(key.isRescheduled());
         Assert.assertNotNull(result.getUuidKeys());
 
         Assert.assertTrue(Iterables.contains(result.getUuidKeys(), uuidKey1));
@@ -80,6 +81,5 @@ public class TestNextBillingDateNotificationKey {
         // Compatibility mode : Although the  uuidKeys is not in the json, we verify the getter return the right result
         Assert.assertNotNull(result.getUuidKeys());
         Assert.assertEquals(result.getUuidKeys().iterator().next().toString(), "a38c363f-b25b-4287-8ebc-55964e116d2f");
-
     }
 }