killbill-aplcache

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
index 6cd53cf..35ee189 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
@@ -676,8 +676,10 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB {
         assertEquals(tags.size(), 1);
     }
 
-    protected void remove_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+
+    protected void remove_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type, final NextEvent...additionalEvents) throws TagDefinitionApiException, TagApiException {
         busHandler.pushExpectedEvent(NextEvent.TAG);
+        busHandler.pushExpectedEvents(additionalEvents);
         tagUserApi.removeTag(id, type, ControlTagType.AUTO_PAY_OFF.getId(), callContext);
         assertListenerStatus();
     }

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
index 5afe70e..007b7bc 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
@@ -89,9 +89,7 @@ public class TestIntegrationWithAutoPayOff extends TestIntegrationBase {
             assertEquals(cur.getBalance(), cur.getChargedAmount());
         }
 
-        busHandler.pushExpectedEvents(NextEvent.PAYMENT);
-        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
+        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT, NextEvent.PAYMENT);
         addDelayBceauseOfLackOfCorrectSynchro();
 
         invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), callContext);
@@ -133,9 +131,7 @@ public class TestIntegrationWithAutoPayOff extends TestIntegrationBase {
         }
 
         paymentPlugin.makeNextPaymentFailWithError();
-        busHandler.pushExpectedEvents(NextEvent.PAYMENT_ERROR);
-        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
+        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT, NextEvent.PAYMENT_ERROR);
         addDelayBceauseOfLackOfCorrectSynchro();
 
         invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), callContext);
@@ -195,9 +191,7 @@ public class TestIntegrationWithAutoPayOff extends TestIntegrationBase {
 
         // NOW SET PLUGIN TO THROW FAILURES
         paymentPlugin.makeNextPaymentFailWithError();
-        busHandler.pushExpectedEvents(NextEvent.PAYMENT_ERROR);
-        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
+        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT, NextEvent.PAYMENT_ERROR);
         addDelayBceauseOfLackOfCorrectSynchro();
 
         invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), callContext);
@@ -231,15 +225,9 @@ public class TestIntegrationWithAutoPayOff extends TestIntegrationBase {
 
         // REMOVE AUTO_PAY_OFF -> WILL SCHEDULE A PAYMENT_RETRY
         paymentPlugin.clear();
-        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
-        assertListenerStatus();
+        remove_AUTO_PAY_OFF_Tag(account.getId(), ObjectType.ACCOUNT, NextEvent.PAYMENT);
         addDelayBceauseOfLackOfCorrectSynchro();
 
-        //
-        busHandler.pushExpectedEvents(NextEvent.PAYMENT);
-        clock.addDays(nbDaysBeforeRetry + 1);
-        assertListenerStatus();
-
         invoices = invoiceUserApi.getInvoicesByAccount(account.getId(), callContext);
         for (Invoice cur : invoices) {
             if (cur.getChargedAmount().compareTo(BigDecimal.ZERO) == 0) {

diff --git a/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java b/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
index 943025c..277389a 100644
--- a/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
+++ b/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
@@ -114,10 +114,6 @@ public final class InvoicePaymentControlPluginApi implements PaymentControlPlugi
         Preconditions.checkArgument(transactionType == TransactionType.PURCHASE ||
                                     transactionType == TransactionType.REFUND);
 
-        if (insert_AUTO_PAY_OFF_ifRequired(paymentControlContext)) {
-            return new DefaultPriorPaymentControlResult(true);
-        }
-
         final InternalCallContext internalContext = internalCallContextFactory.createInternalCallContext(paymentControlContext.getAccountId(), paymentControlContext);
         if (transactionType == TransactionType.PURCHASE) {
             return getPluginPurchaseResult(paymentControlContext, internalContext);
@@ -195,7 +191,12 @@ public final class InvoicePaymentControlPluginApi implements PaymentControlPlugi
             final UUID invoiceId = getInvoiceId(paymentControlPluginContext);
             final Invoice invoice = rebalanceAndGetInvoice(invoiceId, internalContext);
             final BigDecimal requestedAmount = validateAndComputePaymentAmount(invoice, paymentControlPluginContext.getAmount(), paymentControlPluginContext.isApiPayment());
+
             final boolean isAborted = requestedAmount.compareTo(BigDecimal.ZERO) == 0;
+            if (!isAborted && insert_AUTO_PAY_OFF_ifRequired(paymentControlPluginContext, requestedAmount)) {
+                return new DefaultPriorPaymentControlResult(true);
+            }
+
             if (paymentControlPluginContext.isApiPayment() && isAborted) {
                 throw new PaymentControlApiException("Payment for invoice " + invoice.getId() +
                                                      " aborted : invoice balance is = " + invoice.getBalance() +