killbill-memoizeit

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
index d6b01c3..a73f676 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
@@ -31,16 +31,16 @@ import org.joda.time.DateTimeZone;
 import com.ning.billing.ErrorCode;
 import com.ning.billing.account.api.Account;
 import com.ning.billing.account.api.AccountApiException;
+import com.ning.billing.account.api.AccountInternalApi;
+import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.block.BlockingChecker;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
+import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
-import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.util.callcontext.TenantContext;
-import com.ning.billing.account.api.AccountInternalApi;
-import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/block/BlockingChecker.java b/entitlement/src/main/java/com/ning/billing/entitlement/block/BlockingChecker.java
index 5095042..c466378 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/block/BlockingChecker.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/block/BlockingChecker.java
@@ -18,9 +18,9 @@ package com.ning.billing.entitlement.block;
 
 import java.util.UUID;
 
+import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.entitlement.api.Blockable;
 import com.ning.billing.entitlement.api.BlockingApiException;
-import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.entitlement.api.BlockingStateType;
 
 public interface BlockingChecker {

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/block/DefaultBlockingChecker.java b/entitlement/src/main/java/com/ning/billing/entitlement/block/DefaultBlockingChecker.java
index a1ca6d8..a2000a2 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/block/DefaultBlockingChecker.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/block/DefaultBlockingChecker.java
@@ -21,16 +21,16 @@ import java.util.UUID;
 
 import com.ning.billing.ErrorCode;
 import com.ning.billing.account.api.Account;
+import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.entitlement.api.Blockable;
 import com.ning.billing.entitlement.api.BlockingApiException;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.subscription.api.SubscriptionBase;
+import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
-import com.ning.billing.callcontext.InternalTenantContext;
-import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 
 import com.google.inject.Inject;
 

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
index 473d600..060b729 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
@@ -23,8 +23,10 @@ import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.api.BlockingState;
+import com.ning.billing.entitlement.api.EntitlementApiException;
+import com.ning.billing.util.entity.dao.EntityDao;
 
-public interface BlockingStateDao {
+public interface BlockingStateDao extends EntityDao<BlockingStateModelDao, BlockingState, EntitlementApiException> {
 
     /**
      * Returns the current state for that specific service

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateModelDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateModelDao.java
index de4e3da..7727fa5 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateModelDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateModelDao.java
@@ -20,13 +20,13 @@ import java.util.UUID;
 
 import org.joda.time.DateTime;
 
+import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
-import com.ning.billing.callcontext.InternalCallContext;
-import com.ning.billing.util.dao.TableName;
 import com.ning.billing.entity.EntityBase;
-import com.ning.billing.util.entity.dao.EntityModelDao;
 import com.ning.billing.junction.DefaultBlockingState;
+import com.ning.billing.util.dao.TableName;
+import com.ning.billing.util.entity.dao.EntityModelDao;
 
 public class BlockingStateModelDao extends EntityBase implements EntityModelDao<BlockingState>{
 

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
index 6c241ab..466e704 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
@@ -28,12 +28,15 @@ import javax.annotation.Nullable;
 
 import org.skife.jdbi.v2.IDBI;
 
+import com.ning.billing.ErrorCode;
 import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.api.BlockingState;
+import com.ning.billing.entitlement.api.EntitlementApiException;
 import com.ning.billing.util.cache.CacheControllerDispatcher;
 import com.ning.billing.util.dao.NonEntityDao;
+import com.ning.billing.util.entity.dao.EntityDaoBase;
 import com.ning.billing.util.entity.dao.EntitySqlDao;
 import com.ning.billing.util.entity.dao.EntitySqlDaoTransactionWrapper;
 import com.ning.billing.util.entity.dao.EntitySqlDaoTransactionalJdbiWrapper;
@@ -43,7 +46,7 @@ import com.google.common.base.Function;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.Ordering;
 
-public class DefaultBlockingStateDao implements BlockingStateDao {
+public class DefaultBlockingStateDao extends EntityDaoBase<BlockingStateModelDao, BlockingState, EntitlementApiException> implements BlockingStateDao {
 
     // Assume the input is blocking states for a single blockable id
     private static final Ordering<BlockingStateModelDao> BLOCKING_STATE_MODEL_DAO_ORDERING = Ordering.<BlockingStateModelDao>from(new Comparator<BlockingStateModelDao>() {
@@ -60,16 +63,20 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
         }
     });
 
-    private final EntitySqlDaoTransactionalJdbiWrapper transactionalSqlDao;
     private final Clock clock;
 
     public DefaultBlockingStateDao(final IDBI dbi, final Clock clock,
                                    final CacheControllerDispatcher cacheControllerDispatcher, final NonEntityDao nonEntityDao) {
-        this.transactionalSqlDao = new EntitySqlDaoTransactionalJdbiWrapper(dbi, clock, cacheControllerDispatcher, nonEntityDao);
+        super(new EntitySqlDaoTransactionalJdbiWrapper(dbi, clock, cacheControllerDispatcher, nonEntityDao), BlockingStateSqlDao.class);
         this.clock = clock;
     }
 
     @Override
+    protected EntitlementApiException generateAlreadyExistsException(final BlockingStateModelDao blockingStateModelDao, final InternalCallContext context) {
+        return new EntitlementApiException(ErrorCode.ENT_ALREADY_BLOCKED, blockingStateModelDao.getBlockableId());
+    }
+
+    @Override
     public BlockingState getBlockingStateForService(final UUID blockableId, final String serviceName, final InternalTenantContext context) {
         return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<BlockingState>() {
             @Override
@@ -194,7 +201,7 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
 
                 // Create the state, if needed
                 if (!blockingStatesToRemove.contains(newBlockingStateModelDao.getId())) {
-                    sqlDao.create(new BlockingStateModelDao(state, context), context);
+                    sqlDao.create(newBlockingStateModelDao, context);
                 }
 
                 return null;

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
index a35b1f3..5c2967b 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
@@ -45,6 +45,7 @@ import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.util.cache.CacheControllerDispatcher;
 import com.ning.billing.util.dao.NonEntityDao;
+import com.ning.billing.util.entity.Pagination;
 
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
@@ -91,6 +92,46 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
     }
 
     @Override
+    public void create(final BlockingStateModelDao entity, final InternalCallContext context) throws EntitlementApiException {
+        delegate.create(entity, context);
+    }
+
+    @Override
+    public Long getRecordId(final UUID id, final InternalTenantContext context) {
+        return delegate.getRecordId(id, context);
+    }
+
+    @Override
+    public BlockingStateModelDao getByRecordId(final Long recordId, final InternalTenantContext context) {
+        return delegate.getByRecordId(recordId, context);
+    }
+
+    @Override
+    public BlockingStateModelDao getById(final UUID id, final InternalTenantContext context) {
+        return delegate.getById(id, context);
+    }
+
+    @Override
+    public Pagination<BlockingStateModelDao> getAll(final InternalTenantContext context) {
+        return delegate.getAll(context);
+    }
+
+    @Override
+    public Pagination<BlockingStateModelDao> get(final Long offset, final Long limit, final InternalTenantContext context) {
+        return delegate.get(offset, limit, context);
+    }
+
+    @Override
+    public Long getCount(final InternalTenantContext context) {
+        return delegate.getCount(context);
+    }
+
+    @Override
+    public void test(final InternalTenantContext context) {
+        delegate.test(context);
+    }
+
+    @Override
     public BlockingState getBlockingStateForService(final UUID blockableId, final String serviceName, final InternalTenantContext context) {
         return delegate.getBlockingStateForService(blockableId, serviceName, context);
     }

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/DefaultEntitlementService.java b/entitlement/src/main/java/com/ning/billing/entitlement/DefaultEntitlementService.java
index ee87f96..7e7174a 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/DefaultEntitlementService.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/DefaultEntitlementService.java
@@ -27,12 +27,12 @@ import com.ning.billing.bus.api.BusEvent;
 import com.ning.billing.bus.api.PersistentBus;
 import com.ning.billing.bus.api.PersistentBus.EventBusException;
 import com.ning.billing.callcontext.InternalCallContext;
-import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.api.DefaultBlockingTransitionInternalEvent;
 import com.ning.billing.entitlement.api.DefaultEntitlement;
 import com.ning.billing.entitlement.api.Entitlement;
 import com.ning.billing.entitlement.api.EntitlementApi;
 import com.ning.billing.entitlement.api.EntitlementApiException;
+import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.entitlement.engine.core.BlockingTransitionNotificationKey;
 import com.ning.billing.entitlement.engine.core.EntitlementNotificationKey;
 import com.ning.billing.entitlement.engine.core.EntitlementNotificationKeyAction;
@@ -58,8 +58,8 @@ public class DefaultEntitlementService implements EntitlementService {
     private static final Logger log = LoggerFactory.getLogger(DefaultEntitlementService.class);
 
     private final EntitlementApi entitlementApi;
+    private final BlockingStateDao blockingStateDao;
     private final NonEntityDao nonEntityDao;
-    private final Clock clock;
     private final PersistentBus eventBus;
     private final NotificationQueueService notificationQueueService;
     private final InternalCallContextFactory internalCallContextFactory;
@@ -68,14 +68,14 @@ public class DefaultEntitlementService implements EntitlementService {
 
     @Inject
     public DefaultEntitlementService(final EntitlementApi entitlementApi,
+                                     final BlockingStateDao blockingStateDao,
                                      final NonEntityDao nonEntityDao,
-                                     final Clock clock,
                                      final PersistentBus eventBus,
                                      final NotificationQueueService notificationQueueService,
                                      final InternalCallContextFactory internalCallContextFactory) {
         this.entitlementApi = entitlementApi;
+        this.blockingStateDao = blockingStateDao;
         this.nonEntityDao = nonEntityDao;
-        this.clock = clock;
         this.eventBus = eventBus;
         this.notificationQueueService = notificationQueueService;
         this.internalCallContextFactory = internalCallContextFactory;
@@ -141,6 +141,12 @@ public class DefaultEntitlementService implements EntitlementService {
     }
 
     private void processBlockingNotification(final BlockingTransitionNotificationKey key, final InternalCallContext internalCallContext) {
+        // Check if the blocking state has been deleted since
+        if (blockingStateDao.getById(key.getBlockingStateId(), internalCallContext) == null) {
+            log.debug("BlockingState {} has been deleted, not sending a bus event", key.getBlockingStateId());
+            return;
+        }
+
         final BusEvent event = new DefaultBlockingTransitionInternalEvent(key.getBlockableId(), key.getBlockingType(),
                                                                           key.isTransitionedToBlockedBilling(), key.isTransitionedToUnblockedBilling(),
                                                                           key.isTransitionedToBlockedEntitlement(), key.isTransitionToUnblockedEntitlement(),

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/BlockingTransitionNotificationKey.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/BlockingTransitionNotificationKey.java
index 441877e..72ce22d 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/BlockingTransitionNotificationKey.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/BlockingTransitionNotificationKey.java
@@ -26,6 +26,7 @@ import com.fasterxml.jackson.annotation.JsonProperty;
 
 public class BlockingTransitionNotificationKey implements NotificationEvent {
 
+    private final UUID blockingStateId;
     private final UUID blockableId;
     private final BlockingStateType blockingType;
     private final Boolean isTransitionToBlockedBilling;
@@ -34,13 +35,15 @@ public class BlockingTransitionNotificationKey implements NotificationEvent {
     private final Boolean isTransitionToUnblockedEntitlement;
 
     @JsonCreator
-    public BlockingTransitionNotificationKey(@JsonProperty("blockableId") final UUID blockableId,
+    public BlockingTransitionNotificationKey(@JsonProperty("blockingStateId") final UUID blockingStateId,
+                                             @JsonProperty("blockableId") final UUID blockableId,
                                              @JsonProperty("type") final BlockingStateType blockingType,
                                              @JsonProperty("isTransitionToBlockedBilling") final Boolean isTransitionToBlockedBilling,
                                              @JsonProperty("isTransitionToUnblockedBilling") final Boolean isTransitionToUnblockedBilling,
                                              @JsonProperty("isTransitionToBlockedEntitlement") final Boolean isTransitionToBlockedEntitlement,
                                              @JsonProperty("isTransitionToUnblockedEntitlement") final Boolean isTransitionToUnblockedEntitlement) {
 
+        this.blockingStateId = blockingStateId;
         this.blockableId = blockableId;
         this.blockingType = blockingType;
         this.isTransitionToBlockedBilling = isTransitionToBlockedBilling;
@@ -49,6 +52,10 @@ public class BlockingTransitionNotificationKey implements NotificationEvent {
         this.isTransitionToUnblockedEntitlement = isTransitionToUnblockedEntitlement;
     }
 
+    public UUID getBlockingStateId() {
+        return blockingStateId;
+    }
+
     public UUID getBlockableId() {
         return blockableId;
     }
@@ -80,7 +87,8 @@ public class BlockingTransitionNotificationKey implements NotificationEvent {
     @Override
     public String toString() {
         final StringBuilder sb = new StringBuilder("BlockingTransitionNotificationKey{");
-        sb.append("blockableId=").append(blockableId);
+        sb.append("blockingStateId=").append(blockingStateId);
+        sb.append(", blockableId=").append(blockableId);
         sb.append(", blockingType=").append(blockingType);
         sb.append(", isTransitionToBlockedBilling=").append(isTransitionToBlockedBilling);
         sb.append(", isTransitionToUnblockedBilling=").append(isTransitionToUnblockedBilling);
@@ -101,6 +109,9 @@ public class BlockingTransitionNotificationKey implements NotificationEvent {
 
         final BlockingTransitionNotificationKey that = (BlockingTransitionNotificationKey) o;
 
+        if (blockingStateId != null ? !blockingStateId.equals(that.blockingStateId) : that.blockingStateId != null) {
+            return false;
+        }
         if (blockableId != null ? !blockableId.equals(that.blockableId) : that.blockableId != null) {
             return false;
         }
@@ -125,7 +136,8 @@ public class BlockingTransitionNotificationKey implements NotificationEvent {
 
     @Override
     public int hashCode() {
-        int result = blockableId != null ? blockableId.hashCode() : 0;
+        int result = blockingStateId != null ? blockingStateId.hashCode() : 0;
+        result = 31 * result + (blockableId != null ? blockableId.hashCode() : 0);
         result = 31 * result + (blockingType != null ? blockingType.hashCode() : 0);
         result = 31 * result + (isTransitionToBlockedBilling != null ? isTransitionToBlockedBilling.hashCode() : 0);
         result = 31 * result + (isTransitionToUnblockedBilling != null ? isTransitionToUnblockedBilling.hashCode() : 0);

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EntitlementUtils.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EntitlementUtils.java
index 9c440da..c595673 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EntitlementUtils.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EntitlementUtils.java
@@ -77,7 +77,7 @@ public class EntitlementUtils {
 
         final BlockingAggregator currentState = getBlockingStateFor(state.getBlockedId(), state.getType(), context);
         if (previousState != null && currentState != null) {
-            postBlockingTransitionEvent(state.getEffectiveDate(), state.getBlockedId(), state.getType(), previousState, currentState, context);
+            postBlockingTransitionEvent(state.getId(), state.getEffectiveDate(), state.getBlockedId(), state.getType(), previousState, currentState, context);
         }
     }
 
@@ -90,7 +90,7 @@ public class EntitlementUtils {
         }
     }
 
-    private void postBlockingTransitionEvent(final DateTime effectiveDate, final UUID blockableId, final BlockingStateType type,
+    private void postBlockingTransitionEvent(final UUID blockingStateId, final DateTime effectiveDate, final UUID blockableId, final BlockingStateType type,
                                              final BlockingAggregator previousState, final BlockingAggregator currentState,
                                              final InternalCallContext context) {
         final boolean isTransitionToBlockedBilling = !previousState.isBlockBilling() && currentState.isBlockBilling();
@@ -101,7 +101,7 @@ public class EntitlementUtils {
 
         if (effectiveDate.compareTo(clock.getUTCNow()) > 0) {
             // Add notification entry to send the bus event at the effective date
-            final NotificationEvent notificationEvent = new BlockingTransitionNotificationKey(blockableId, type,
+            final NotificationEvent notificationEvent = new BlockingTransitionNotificationKey(blockingStateId, blockableId, type,
                                                                                               isTransitionToBlockedBilling, isTransitionToUnblockedBilling,
                                                                                               isTransitionToBlockedEntitlement, isTransitionToUnblockedEntitlement);
             recordFutureNotification(effectiveDate, notificationEvent, context);

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlement.java b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlement.java
index 5778669..93cef92 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlement.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlement.java
@@ -32,7 +32,6 @@ import com.ning.billing.entitlement.api.Entitlement.EntitlementActionPolicy;
 import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 
 import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertTrue;
 
 public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
 
@@ -48,7 +47,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement.getState(), EntitlementState.ACTIVE);
 
         clock.addDays(5);
@@ -56,7 +55,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
         final LocalDate cancelDate = new LocalDate(clock.getUTCNow());
         entitlement.cancelEntitlementWithDate(cancelDate, true, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement2 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
         assertEquals(entitlement2.getState(), EntitlementState.CANCELLED);
@@ -75,7 +74,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement.getState(), EntitlementState.ACTIVE);
 
         clock.addDays(5);
@@ -90,7 +89,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         clock.addDays(1);
 
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement3 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
         assertEquals(entitlement3.getState(), EntitlementState.CANCELLED);
@@ -109,7 +108,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement.getState(), EntitlementState.ACTIVE);
 
         clock.addDays(5);
@@ -123,7 +122,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
 
         testListener.pushExpectedEvents(NextEvent.UNCANCEL);
         entitlement2.uncancelEntitlement(callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         clock.addDays(1);
         final Entitlement entitlement3 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
@@ -142,11 +141,11 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
         final Entitlement cancelledEntitlement = entitlement.cancelEntitlementWithPolicy(EntitlementActionPolicy.END_OF_TERM, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(cancelledEntitlement.getState(), EntitlementState.CANCELLED);
         assertEquals(cancelledEntitlement.getEffectiveEndDate(), initialDate);
 
@@ -167,7 +166,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final DateTime ctd = clock.getUTCNow().plusDays(30).plusMonths(1);
         testListener.pushExpectedEvent(NextEvent.PHASE);
@@ -175,13 +174,13 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         clock.addDays(32);
         // Set manually since no invoice
         subscriptionInternalApi.setChargedThroughDate(entitlement.getId(), ctd, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement2 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
 
         testListener.pushExpectedEvent(NextEvent.BLOCK);
         final Entitlement entitlement3 = entitlement2.cancelEntitlementWithPolicy(EntitlementActionPolicy.IMMEDIATE, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement3.getState(), EntitlementState.CANCELLED);
         assertEquals(entitlement3.getEffectiveEndDate(), new LocalDate(2013, 9, 8));
 
@@ -191,7 +190,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
 
         testListener.pushExpectedEvent(NextEvent.CANCEL);
         clock.addMonths(1);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement4 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
         assertEquals(entitlement4.getState(), EntitlementState.CANCELLED);
@@ -210,14 +209,14 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final DateTime ctd = clock.getUTCNow().plusDays(30).plusMonths(1);
         testListener.pushExpectedEvent(NextEvent.PHASE);
         clock.addDays(32);
         // Set manually since no invoice
         subscriptionInternalApi.setChargedThroughDate(entitlement.getId(), ctd, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement2 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
 
@@ -231,7 +230,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
 
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
         clock.addMonths(1);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement entitlement4 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
         assertEquals(entitlement4.getState(), EntitlementState.CANCELLED);
@@ -250,12 +249,12 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Immediate change during trial
         testListener.pushExpectedEvent(NextEvent.CHANGE);
         entitlement.changePlan("Assault-Rifle", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Verify the change is immediate
         final Entitlement entitlement2 = entitlementApi.getEntitlementForId(entitlement.getId(), callContext);
@@ -263,7 +262,7 @@ public class TestDefaultEntitlement extends EntitlementTestSuiteWithEmbeddedDB {
 
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
         final Entitlement cancelledEntitlement = entitlement.cancelEntitlementWithPolicy(EntitlementActionPolicy.END_OF_TERM, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(cancelledEntitlement.getState(), EntitlementState.CANCELLED);
         assertEquals(cancelledEntitlement.getEffectiveEndDate(), initialDate);
 

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlementApi.java b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlementApi.java
index 612a949..b2a4067 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlementApi.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultEntitlementApi.java
@@ -39,7 +39,6 @@ import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNull;
-import static org.testng.Assert.assertTrue;
 
 
 public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedDB {
@@ -56,14 +55,14 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Keep the same object for the whole test, to make sure we refresh its state before r/w calls
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Add ADD_ON
         // Keep the same object for the whole test, to make sure we refresh its state before r/w calls
         final PlanPhaseSpecifier addOnSpec = new PlanPhaseSpecifier("Telescopic-Scope", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement addOnEntitlement = entitlementApi.addEntitlement(entitlement.getBundleId(), addOnSpec, initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         /*
         // TODO It looks like we don't check if there is a future cancellation. Maybe we should?
@@ -79,7 +78,7 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Cancelling the base entitlement will cancel the add-on
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK, NextEvent.CANCEL, NextEvent.BLOCK);
         entitlement.cancelEntitlementWithDateOverrideBillingPolicy(clock.getUTCToday(), BillingActionPolicy.IMMEDIATE, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         try {
             entitlement.cancelEntitlementWithDateOverrideBillingPolicy(clock.getUTCToday(), BillingActionPolicy.IMMEDIATE, callContext);
@@ -122,7 +121,7 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement.getAccountId(), account.getId());
         assertEquals(entitlement.getExternalKey(), account.getExternalKey());
 
@@ -214,13 +213,13 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement baseEntitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Add ADD_ON
         final PlanPhaseSpecifier spec1 = new PlanPhaseSpecifier("Telescopic-Scope", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement telescopicEntitlement = entitlementApi.addEntitlement(baseEntitlement.getBundleId(), spec1, initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         assertEquals(telescopicEntitlement.getAccountId(), account.getId());
         assertEquals(telescopicEntitlement.getExternalKey(), account.getExternalKey());
@@ -253,14 +252,14 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement baseEntitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         clock.addDays(1);
         final LocalDate effectiveDateSpec1 = new LocalDate(clock.getUTCNow(), account.getTimeZone());
         final PlanPhaseSpecifier spec1 = new PlanPhaseSpecifier("Telescopic-Scope", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement telescopicEntitlement = entitlementApi.addEntitlement(baseEntitlement.getBundleId(), spec1, effectiveDateSpec1, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Block all entitlement in the bundle
         clock.addDays(5);
@@ -337,20 +336,20 @@ public class TestDefaultEntitlementApi extends EntitlementTestSuiteWithEmbeddedD
         // Create entitlement
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement baseEntitlement = entitlementApi.createBaseEntitlement(accountSrc.getId(), spec, accountSrc.getExternalKey(), initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final DateTime ctd = clock.getUTCNow().plusDays(30).plusMonths(1);
         testListener.pushExpectedEvent(NextEvent.PHASE);
         clock.addDays(32);
         // Set manually since no invoice
         subscriptionInternalApi.setChargedThroughDate(baseEntitlement.getId(), ctd, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Transfer bundle to dest account
         final LocalDate effectiveDate = new LocalDate(clock.getUTCNow(), accountSrc.getTimeZone());
         testListener.pushExpectedEvents(NextEvent.TRANSFER, NextEvent.BLOCK);
         final UUID newBundleId = entitlementApi.transferEntitlementsOverrideBillingPolicy(accountSrc.getId(), accountDesc.getId(), baseEntitlement.getExternalKey(), effectiveDate, BillingActionPolicy.END_OF_TERM, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final Entitlement oldBaseEntitlement = entitlementApi.getAllEntitlementsForAccountIdAndExternalKey(accountSrc.getId(), accountSrc.getExternalKey(), callContext).get(0);
         assertEquals(oldBaseEntitlement.getEffectiveEndDate(), effectiveDate);

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
index 0ffe7ba..ea79a11 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
@@ -32,7 +32,6 @@ import com.ning.billing.entitlement.EntitlementTestSuiteWithEmbeddedDB;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNull;
-import static org.testng.Assert.assertTrue;
 
 public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbeddedDB {
 
@@ -51,7 +50,7 @@ public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbedded
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, externalKey, initialDate, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement.getAccountId(), account.getId());
         assertEquals(entitlement.getExternalKey(), externalKey);
 
@@ -65,7 +64,7 @@ public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbedded
         clock.addDays(3);
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK);
         entitlement.cancelEntitlementWithDate(new LocalDate(clock.getUTCNow(), account.getTimeZone()), true, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         clock.addDays(1);
         // Re-create a new bundle with same externalKey
@@ -74,7 +73,7 @@ public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbedded
         // Create entitlement and check each field
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement2 = entitlementApi.createBaseEntitlement(account.getId(), spec2, externalKey, new LocalDate(clock.getUTCNow(), account.getTimeZone()), callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
         assertEquals(entitlement2.getAccountId(), account.getId());
         assertEquals(entitlement2.getExternalKey(), externalKey);
 
@@ -105,7 +104,7 @@ public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbedded
 
         testListener.pushExpectedEvents(NextEvent.TRANSFER, NextEvent.CANCEL, NextEvent.BLOCK);
         entitlementApi.transferEntitlements(account.getId(), account2.getId(), externalKey, new LocalDate(clock.getUTCNow(), account.getTimeZone()), callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final List<SubscriptionBundle> bundles3 = subscriptionApi.getSubscriptionBundlesForExternalKey(externalKey, callContext);
         assertEquals(bundles3.size(), 3);

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestEntitlementDateHelper.java b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestEntitlementDateHelper.java
index 68576be..d1af0b3 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestEntitlementDateHelper.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestEntitlementDateHelper.java
@@ -25,8 +25,8 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import com.ning.billing.account.api.Account;
-import com.ning.billing.entitlement.EntitlementTestSuiteNoDB;
 import com.ning.billing.callcontext.InternalTenantContext;
+import com.ning.billing.entitlement.EntitlementTestSuiteNoDB;
 
 import static org.testng.Assert.assertTrue;
 

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingApi.java b/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingApi.java
index 395e22e..edd57ed 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingApi.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingApi.java
@@ -28,10 +28,8 @@ import com.ning.billing.api.TestApiListener.NextEvent;
 import com.ning.billing.entitlement.EntitlementTestSuiteWithEmbeddedDB;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
-import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.junction.DefaultBlockingState;
-
-import static org.testng.Assert.assertTrue;
+import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 
 public class TestBlockingApi extends EntitlementTestSuiteWithEmbeddedDB {
 
@@ -52,9 +50,9 @@ public class TestBlockingApi extends EntitlementTestSuiteWithEmbeddedDB {
         final boolean blockBilling = false;
 
         testListener.pushExpectedEvent(NextEvent.BLOCK);
-        final BlockingState state1 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT,overdueStateName, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
+        final BlockingState state1 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
         blockingInternalApi.setBlockingState(state1, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         clock.setDeltaFromReality(1000 * 3600 * 24);
 
@@ -62,7 +60,7 @@ public class TestBlockingApi extends EntitlementTestSuiteWithEmbeddedDB {
         final String overdueStateName2 = "NoReallyThisCantGoOn";
         final BlockingState state2 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName2, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
         blockingInternalApi.setBlockingState(state2, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final SubscriptionBaseBundle bundle = Mockito.mock(SubscriptionBaseBundle.class);
         Mockito.when(bundle.getId()).thenReturn(uuid);
@@ -84,7 +82,7 @@ public class TestBlockingApi extends EntitlementTestSuiteWithEmbeddedDB {
         testListener.pushExpectedEvent(NextEvent.BLOCK);
         final BlockingState state1 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
         blockingInternalApi.setBlockingState(state1, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         clock.setDeltaFromReality(1000 * 3600 * 24);
 
@@ -92,7 +90,7 @@ public class TestBlockingApi extends EntitlementTestSuiteWithEmbeddedDB {
         final String overdueStateName2 = "NoReallyThisCantGoOn";
         final BlockingState state2 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName2, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
         blockingInternalApi.setBlockingState(state2, internalCallContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         final SubscriptionBaseBundle bundle = Mockito.mock(SubscriptionBaseBundle.class);
         Mockito.when(bundle.getId()).thenReturn(uuid);

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingChecker.java b/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingChecker.java
index eed0ce0..469453d 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingChecker.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/block/TestBlockingChecker.java
@@ -17,21 +17,23 @@
 package com.ning.billing.entitlement.block;
 
 import java.util.UUID;
+
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
 import com.ning.billing.account.api.Account;
+import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.entitlement.EntitlementTestSuiteNoDB;
-import com.ning.billing.entitlement.api.BlockingStateType;
-import com.ning.billing.entitlement.dao.MockBlockingStateDao;
 import com.ning.billing.entitlement.api.BlockingApiException;
 import com.ning.billing.entitlement.api.BlockingState;
+import com.ning.billing.entitlement.api.BlockingStateType;
+import com.ning.billing.entitlement.dao.MockBlockingStateDao;
+import com.ning.billing.junction.DefaultBlockingState;
 import com.ning.billing.subscription.api.SubscriptionBase;
-import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
-import com.ning.billing.callcontext.InternalTenantContext;
-import com.ning.billing.junction.DefaultBlockingState;
-import org.mockito.Mockito;
-import org.testng.Assert;
-import org.testng.annotations.BeforeMethod;
-import org.testng.annotations.Test;
+import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 
 public class TestBlockingChecker extends EntitlementTestSuiteNoDB {
 

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java b/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
index 5dae30d..fa343d8 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
@@ -24,16 +24,18 @@ import java.util.UUID;
 
 import javax.annotation.Nullable;
 
-import com.ning.billing.clock.Clock;
-import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.callcontext.InternalTenantContext;
+import com.ning.billing.clock.Clock;
+import com.ning.billing.entitlement.api.BlockingState;
+import com.ning.billing.entitlement.api.EntitlementApiException;
+import com.ning.billing.util.entity.dao.MockEntityDaoBase;
 
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
 
-public class MockBlockingStateDao implements BlockingStateDao {
+public class MockBlockingStateDao extends MockEntityDaoBase<BlockingStateModelDao, BlockingState, EntitlementApiException> implements BlockingStateDao {
 
     private final Map<UUID, List<BlockingState>> blockingStates = new HashMap<UUID, List<BlockingState>>();
 
@@ -61,7 +63,7 @@ public class MockBlockingStateDao implements BlockingStateDao {
             return new ArrayList<BlockingState>();
         }
 
-        final Map<String, BlockingState> tmp  = new HashMap<String, BlockingState>();
+        final Map<String, BlockingState> tmp = new HashMap<String, BlockingState>();
         for (BlockingState cur : blockingStatesForId) {
             final BlockingState curStateForService = tmp.get(cur.getService());
             if (curStateForService == null || curStateForService.getEffectiveDate().compareTo(cur.getEffectiveDate()) < 0) {

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
index c9d3e53..7850378 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
@@ -27,8 +27,8 @@ import org.testng.annotations.Test;
 import com.ning.billing.entitlement.EntitlementTestSuiteWithEmbeddedDB;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
-import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.junction.DefaultBlockingState;
+import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 
 public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
 

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java b/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
index 8a066aa..abef851 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
@@ -45,8 +45,6 @@ import com.ning.billing.entitlement.api.Entitlement.EntitlementActionPolicy;
 import com.ning.billing.entitlement.api.EntitlementApiException;
 import com.ning.billing.entitlement.dao.BlockingStateSqlDao;
 
-import static org.testng.Assert.assertTrue;
-
 public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
 
     private BlockingStateSqlDao sqlDao;
@@ -82,7 +80,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         testListener.pushExpectedEvents(NextEvent.PHASE, NextEvent.PHASE);
         // Phase for the base plan is 2013/09/07 (30 days trial) but it's 2013/09/08 for the add-on (1 month discount)
         clock.setDay(new LocalDate(2013, 9, 8));
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Note! Make sure to align CTD and cancellation/change effective time with the phase event effective time to avoid timing issues in comparisons
         baseEffectiveEOTCancellationOrChangeDateTime = baseEntitlement.getSubscriptionBase().getAllTransitions().get(1).getEffectiveTransitionTime().plusMonths(1);
@@ -97,7 +95,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Cancel the base plan
         final DefaultEntitlement cancelledBaseEntitlement = (DefaultEntitlement) baseEntitlement.cancelEntitlementWithPolicyOverrideBillingPolicy(EntitlementActionPolicy.END_OF_TERM, BillingActionPolicy.END_OF_TERM, callContext);
         // No blocking event (EOT)
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Verify we compute the right blocking states for the "read" path...
         checkFutureBlockingStatesToCancel(addOnEntitlement, null, null);
@@ -111,7 +109,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Verify the notification kicks in
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK, NextEvent.CANCEL, NextEvent.BLOCK);
         clock.addDays(30);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Refresh the state
         final DefaultEntitlement cancelledAddOnEntitlement = (DefaultEntitlement) entitlementApi.getEntitlementForId(addOnEntitlement.getId(), callContext);
@@ -135,7 +133,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Cancel the base plan
         testListener.pushExpectedEvents(NextEvent.CANCEL, NextEvent.BLOCK, NextEvent.CANCEL, NextEvent.BLOCK);
         final DefaultEntitlement cancelledBaseEntitlement = (DefaultEntitlement) baseEntitlement.cancelEntitlementWithPolicyOverrideBillingPolicy(EntitlementActionPolicy.IMMEDIATE, BillingActionPolicy.IMMEDIATE, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Refresh the add-on state
         final DefaultEntitlement cancelledAddOnEntitlement = (DefaultEntitlement) entitlementApi.getEntitlementForId(addOnEntitlement.getId(), callContext);
@@ -151,7 +149,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
 
         clock.addDays(30);
         // No new event
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
@@ -165,7 +163,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Change plan EOT to Assault-Rifle (Telescopic-Scope is included)
         final DefaultEntitlement changedBaseEntitlement = (DefaultEntitlement) baseEntitlement.changePlanWithDate("Assault-Rifle", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, new LocalDate(2013, 10, 7), callContext);
         // No blocking event (EOT)
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Verify we compute the right blocking states for the "read" path...
         checkFutureBlockingStatesToCancel(addOnEntitlement, null, null);
@@ -179,7 +177,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Verify the notification kicks in
         testListener.pushExpectedEvents(NextEvent.CHANGE, NextEvent.CANCEL, NextEvent.BLOCK);
         clock.addDays(30);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Refresh the state
         final DefaultEntitlement cancelledAddOnEntitlement = (DefaultEntitlement) entitlementApi.getEntitlementForId(addOnEntitlement.getId(), callContext);
@@ -203,7 +201,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Change plan IMM (upgrade) to Assault-Rifle (Telescopic-Scope is included)
         testListener.pushExpectedEvents(NextEvent.CHANGE, NextEvent.CANCEL, NextEvent.BLOCK);
         final DefaultEntitlement changedBaseEntitlement = (DefaultEntitlement) baseEntitlement.changePlan("Assault-Rifle", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, callContext);
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         // Refresh the add-on state
         final DefaultEntitlement cancelledAddOnEntitlement = (DefaultEntitlement) entitlementApi.getEntitlementForId(addOnEntitlement.getId(), callContext);
@@ -219,7 +217,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
 
         clock.addDays(30);
         // No new event
-        assertTrue(testListener.isCompleted(DELAY));
+        assertListenerStatus();
 
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteNoDB.java b/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteNoDB.java
index 8bf1fff..b1428fd 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteNoDB.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteNoDB.java
@@ -21,12 +21,12 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 
 import com.ning.billing.GuicyKillbillTestSuiteNoDB;
+import com.ning.billing.account.api.AccountInternalApi;
 import com.ning.billing.bus.api.PersistentBus;
 import com.ning.billing.catalog.api.CatalogService;
 import com.ning.billing.entitlement.block.BlockingChecker;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.entitlement.glue.TestEntitlementModuleNoDB;
-import com.ning.billing.account.api.AccountInternalApi;
 import com.ning.billing.junction.BlockingInternalApi;
 import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
 import com.ning.billing.tag.TagInternalApi;
@@ -54,9 +54,7 @@ public abstract class EntitlementTestSuiteNoDB extends GuicyKillbillTestSuiteNoD
     protected TagDao tagDao;
     @Inject
     protected TagInternalApi tagInternalApi;
-
     @Inject
-
     protected BlockingChecker blockingChecker;
 
     @BeforeClass(groups = "fast")
@@ -74,6 +72,4 @@ public abstract class EntitlementTestSuiteNoDB extends GuicyKillbillTestSuiteNoD
     public void afterMethod() {
         bus.stop();
     }
-
-
 }

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java b/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java
index 9f7c6ec..d4c2131 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java
@@ -61,6 +61,7 @@ import com.google.inject.Injector;
 import com.google.inject.Stage;
 
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
 
 public class EntitlementTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuiteWithEmbeddedDB {
 
@@ -235,4 +236,9 @@ public class EntitlementTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuiteWi
                                        .timeZone(DateTimeZone.UTC)
                                        .build();
     }
+
+    protected void assertListenerStatus() {
+        assertTrue(testListener.isCompleted(DELAY));
+        ((EntitlementTestListenerStatus) testListenerStatus).assertListenerStatus();
+    }
 }

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModule.java b/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModule.java
index 7306110..a221909 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModule.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModule.java
@@ -16,13 +16,10 @@
 
 package com.ning.billing.entitlement.glue;
 
-import com.ning.billing.catalog.MockCatalogModule;
-import com.ning.billing.mock.glue.MockAccountModule;
-import com.ning.billing.mock.glue.MockSubscriptionModule;
-import com.ning.billing.subscription.glue.DefaultSubscriptionModule;
+import org.skife.config.ConfigSource;
+
 import com.ning.billing.util.glue.CacheModule;
 import com.ning.billing.util.glue.CallContextModule;
-import org.skife.config.ConfigSource;
 
 public class TestEntitlementModule extends DefaultEntitlementModule {
 

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModuleWithEmbeddedDB.java b/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModuleWithEmbeddedDB.java
index b104bc5..2369cb6 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModuleWithEmbeddedDB.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/glue/TestEntitlementModuleWithEmbeddedDB.java
@@ -16,20 +16,20 @@
 
 package com.ning.billing.entitlement.glue;
 
+import org.skife.config.ConfigSource;
+
 import com.ning.billing.GuicyKillbillTestWithEmbeddedDBModule;
 import com.ning.billing.account.glue.DefaultAccountModule;
 import com.ning.billing.api.TestApiListener;
 import com.ning.billing.api.TestListenerStatus;
 import com.ning.billing.catalog.glue.CatalogModule;
 import com.ning.billing.entitlement.EntitlementTestListenerStatus;
-import com.ning.billing.glue.AccountModule;
 import com.ning.billing.subscription.glue.DefaultSubscriptionModule;
 import com.ning.billing.util.glue.BusModule;
 import com.ning.billing.util.glue.MetricsModule;
 import com.ning.billing.util.glue.NonEntityDaoModule;
 import com.ning.billing.util.glue.NotificationQueueModule;
 import com.ning.billing.util.glue.TagStoreModule;
-import org.skife.config.ConfigSource;
 
 public class TestEntitlementModuleWithEmbeddedDB extends TestEntitlementModule {