killbill-memoizeit

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
index c2d7df1..714ba92 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
@@ -721,7 +721,7 @@ public class TestOverdueIntegration extends TestOverdueBase {
 
         // Pay the invoice
         final Invoice externalChargeInvoice = invoiceUserApi.getUnpaidInvoicesByAccountId(account.getId(), clock.getUTCToday(), callContext).iterator().next();
-        createExternalPaymentAndCheckForCompletion(account, externalChargeInvoice, NextEvent.PAYMENT);
+        createExternalPaymentAndCheckForCompletion(account, externalChargeInvoice, NextEvent.PAYMENT, NextEvent.BLOCK);
         // We should be clear now
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
     }
@@ -853,7 +853,7 @@ public class TestOverdueIntegration extends TestOverdueBase {
         // Pay the first one via an external payment - we should then be 5 days apart from the second invoice
         // (which is the earliest unpaid one) and hence come back to a clear state (see configuration)
         final Invoice firstNonZeroInvoice = invoiceUserApi.getUnpaidInvoicesByAccountId(account.getId(), clock.getUTCToday(), callContext).iterator().next();
-        createExternalPaymentAndCheckForCompletion(account, firstNonZeroInvoice, NextEvent.PAYMENT);
+        createExternalPaymentAndCheckForCompletion(account, firstNonZeroInvoice, NextEvent.PAYMENT, NextEvent.BLOCK);
         // We should be clear now
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
     }

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueWithOverdueEnforcementOffTag.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueWithOverdueEnforcementOffTag.java
index afda05d..a9694a2 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueWithOverdueEnforcementOffTag.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueWithOverdueEnforcementOffTag.java
@@ -68,7 +68,7 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         // Set the OVERDUE_ENFORCEMENT_OFF tag (we set the clear state, hence the blocking event)
         busHandler.pushExpectedEvents(NextEvent.TAG, NextEvent.BLOCK);
         tagUserApi.addTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-        busHandler.isCompleted(DELAY);
+        assertListenerStatus();
 
         // Set next invoice to fail and create subscription
         paymentPlugin.makeAllInvoicesFailWithError(true);
@@ -91,9 +91,9 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
 
         // Now remove OVERDUE_ENFORCEMENT_OFF tag
-        busHandler.pushExpectedEvent(NextEvent.TAG);
+        busHandler.pushExpectedEvents(NextEvent.TAG, NextEvent.BLOCK);
         tagUserApi.removeTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-        busHandler.isCompleted(DELAY);
+        assertListenerStatus();
         checkODState("OD1");
     }
 
@@ -118,23 +118,25 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         invoiceChecker.checkChargedThroughDate(baseEntitlement.getId(), new LocalDate(2012, 6, 30), callContext);
 
         // DAY 36 -- RIGHT AFTER OD1
+        busHandler.pushExpectedEvent(NextEvent.BLOCK);
         addDaysAndCheckForCompletion(6);
+        assertListenerStatus();
 
         // Account should be in overdue
         checkODState("OD1");
 
         // Set the OVERDUE_ENFORCEMENT_OFF tag
-        busHandler.pushExpectedEvent(NextEvent.TAG);
+        busHandler.pushExpectedEvents(NextEvent.TAG, NextEvent.BLOCK);
         tagUserApi.addTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-        busHandler.isCompleted(DELAY);
+        assertListenerStatus();
 
         // Should now be in clear state
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
 
         // Now remove OVERDUE_ENFORCEMENT_OFF tag
-        busHandler.pushExpectedEvent(NextEvent.TAG);
+        busHandler.pushExpectedEvents(NextEvent.TAG,  NextEvent.BLOCK);
         tagUserApi.removeTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-        busHandler.isCompleted(DELAY);
+        assertListenerStatus();
 
         // Account should be back in overdue
         checkODState("OD1");

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
index 3029437..fe2e603 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
@@ -273,10 +273,16 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
         lifecycle.fireStartupSequencePostEventRegistration();
 
         paymentPlugin.clear();
+
+        // Make sure we start with a clean state
+        assertListenerStatus();
     }
 
     @AfterMethod(groups = "slow")
     public void afterMethod() throws Exception {
+        // Make sure we finish in a clean state
+        assertListenerStatus();
+
         lifecycle.fireShutdownSequencePriorEventUnRegistration();
         busService.getBus().unregister(busHandler);
         lifecycle.fireShutdownSequencePostEventUnRegistration();
@@ -610,7 +616,6 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
         busHandler.pushExpectedEvents(events);
 
         final T result = f.apply(null);
-        assertTrue(busHandler.isCompleted(DELAY), "Were expecting events " + joiner.join(events));
         assertListenerStatus();