killbill-memoizeit

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 095cef7..c3bc728 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
@@ -69,7 +69,9 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         clock.setTime(new DateTime(2012, 5, 1, 0, 3, 42, 0));
 
         // Set the OVERDUE_ENFORCEMENT_OFF tag
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagUserApi.addTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
+        busHandler.isCompleted(DELAY);
 
         // Set next invoice to fail and create subscription
         paymentPlugin.makeAllInvoicesFailWithError(true);
@@ -91,8 +93,9 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
 
         // Now remove OVERDUE_ENFORCEMENT_OFF tag
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagUserApi.removeTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-
+        busHandler.isCompleted(DELAY);
         checkODState("OD1");
     }
 
@@ -123,14 +126,18 @@ public class TestOverdueWithOverdueEnforcementOffTag extends TestOverdueBase {
         checkODState("OD1");
 
         // Set the OVERDUE_ENFORCEMENT_OFF tag
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagUserApi.addTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
+        busHandler.isCompleted(DELAY);
 
         // Should now be in clear state
         checkODState(DefaultBlockingState.CLEAR_STATE_NAME);
 
         // Now remove OVERDUE_ENFORCEMENT_OFF tag
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagUserApi.removeTag(account.getId(), ObjectType.ACCOUNT, ControlTagType.OVERDUE_ENFORCEMENT_OFF.getId(), callContext);
-
+        busHandler.isCompleted(DELAY);
+        
         // Account should be back in overdue
         checkODState("OD1");
     }

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
index b73db5d..d88502f 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoPayOff.java
@@ -279,13 +279,17 @@ public class TestIntegrationWithAutoPayOff extends TestIntegrationBase {
 
 
     private void add_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagApi.addTag(id, type, ControlTagType.AUTO_PAY_OFF.getId(), callContext);
+        assertTrue(busHandler.isCompleted(DELAY));
         final List<Tag> tags = tagApi.getTags(id, type, callContext);
         assertEquals(tags.size(), 1);
     }
 
     private void remove_AUTO_PAY_OFF_Tag(final UUID id, final ObjectType type) throws TagDefinitionApiException, TagApiException {
+        busHandler.pushExpectedEvent(NextEvent.TAG);
         tagApi.removeTag(id, type, ControlTagType.AUTO_PAY_OFF.getId(), callContext);
+        assertTrue(busHandler.isCompleted(DELAY));
     }
 }