killbill-aplcache

entitlement: move the add-ons cancellation computation to

11/11/2013 12:47:48 PM

Details

diff --git a/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBase.java b/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBase.java
index a75c5f8..09fabfc 100644
--- a/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBase.java
+++ b/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBase.java
@@ -53,15 +53,15 @@ public interface SubscriptionBase extends Entity, Blockable {
     public boolean uncancel(final CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlan(final String productName, final BillingPeriod term, final String priceList, final CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlanWithDate(final String productName, final BillingPeriod term, final String priceList, final DateTime requestedDate, final CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlanWithPolicy(final String productName, final BillingPeriod term, final String priceList,
                                          final BillingActionPolicy policy, final CallContext context)
             throws SubscriptionBaseApiException;
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlement.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlement.java
index ddca8b5..85d1a1c 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlement.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlement.java
@@ -20,8 +20,6 @@ import java.io.IOException;
 import java.util.Collection;
 import java.util.UUID;
 
-import javax.annotation.Nullable;
-
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 
@@ -318,15 +316,20 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
             throw new EntitlementApiException(ErrorCode.SUB_CANCEL_BAD_STATE, getId(), EntitlementState.CANCELLED);
         }
 
+        // Make sure to compute the entitlement effective date first to avoid timing issues for IMM cancellations
+        // (we don't want an entitlement cancel date one second or so after the subscription cancel date or add-ons cancellations
+        // computations won't work).
+        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
+        final LocalDate effectiveLocalDate = new LocalDate(localCancelDate, eventsStream.getAccount().getTimeZone());
+        final DateTime effectiveDate = dateHelper.fromLocalDateAndReferenceTime(effectiveLocalDate, getSubscriptionBase().getStartDate(), contextWithValidAccountRecordId);
+
         try {
+            // Cancel subscription base first, to correctly compute the add-ons entitlements we need to cancel (see below)
             getSubscriptionBase().cancelWithPolicy(billingPolicy, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
-        final LocalDate effectiveLocalDate = new LocalDate(localCancelDate, eventsStream.getAccount().getTimeZone());
-        final DateTime effectiveDate = dateHelper.fromLocalDateAndReferenceTime(effectiveLocalDate, getSubscriptionBase().getStartDate(), contextWithValidAccountRecordId);
         final BlockingState newBlockingState = new DefaultBlockingState(getId(), BlockingStateType.SUBSCRIPTION, DefaultEntitlementApi.ENT_STATE_CANCELLED, EntitlementService.ENTITLEMENT_SERVICE_NAME, true, true, false, effectiveDate);
         entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(newBlockingState, contextWithValidAccountRecordId);
 
@@ -438,7 +441,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         eventsStream = eventsStreamBuilder.refresh(eventsStream, context);
     }
 
-    public void blockAddOnsIfRequired(@Nullable final DateTime effectiveDateOrNull, final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
+    public void blockAddOnsIfRequired(final DateTime effectiveDate, final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
         // Optimization - bail early
         if (!ProductCategory.BASE.equals(getSubscriptionBase().getCategory())) {
             // Only base subscriptions have add-ons
@@ -448,27 +451,22 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         // Get the latest state from disk (we just got cancelled or changed plan)
         refresh(context);
 
-        final DateTime now = clock.getUTCNow();
-
-        // null means immediate
-        final DateTime effectiveDate = effectiveDateOrNull == null ? now : effectiveDateOrNull;
-
-        final boolean isBaseEntitlementCancelled = eventsStream.isEntitlementCancelled();
-
         // If cancellation/change occurs in the future, do nothing for now but add a notification entry.
         // This is to distinguish whether a future cancellation was requested by the user, or was a side effect
         // (e.g. base plan cancellation): future entitlement cancellations for add-ons on disk always reflect
         // an explicit cancellation. This trick lets us determine what to do when un-cancelling.
         // This mirror the behavior in subscription base (see DefaultSubscriptionBaseApiService).
+        final DateTime now = clock.getUTCNow();
         if (effectiveDate.compareTo(now) > 0) {
             // Note that usually we record the notification from the DAO. We cannot do it here because not all calls
             // go through the DAO (e.g. change)
+            final boolean isBaseEntitlementCancelled = eventsStream.isEntitlementCancelled();
             final NotificationEvent notificationEvent = new EntitlementNotificationKey(getId(), isBaseEntitlementCancelled ? EntitlementNotificationKeyAction.CANCEL : EntitlementNotificationKeyAction.CHANGE, effectiveDate);
             recordFutureNotification(effectiveDate, notificationEvent, internalCallContext);
             return;
         }
 
-        final Collection<BlockingState> addOnsBlockingStates = entitlementUtils.computeBlockingStatesForAssociatedAddons(getSubscriptionBase(), effectiveDate, internalCallContext);
+        final Collection<BlockingState> addOnsBlockingStates = eventsStream.computeAddonsBlockingStatesForNextSubscriptionBaseEvent(effectiveDate);
         for (final BlockingState addOnBlockingState : addOnsBlockingStates) {
             entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(addOnBlockingState, internalCallContext);
         }
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 7ea44e1..6c241ab 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
@@ -21,30 +21,17 @@ import java.util.Collection;
 import java.util.Comparator;
 import java.util.Date;
 import java.util.HashSet;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
-import javax.inject.Inject;
 
-import org.joda.time.DateTime;
 import org.skife.jdbi.v2.IDBI;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.callcontext.InternalTenantContext;
-import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.clock.Clock;
-import com.ning.billing.entitlement.EntitlementService;
 import com.ning.billing.entitlement.api.BlockingState;
-import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
-import com.ning.billing.entitlement.api.EntitlementApiException;
-import com.ning.billing.entitlement.engine.core.EntitlementUtils;
-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.util.cache.CacheControllerDispatcher;
 import com.ning.billing.util.dao.NonEntityDao;
 import com.ning.billing.util.entity.dao.EntitySqlDao;
@@ -53,14 +40,12 @@ import com.ning.billing.util.entity.dao.EntitySqlDaoTransactionalJdbiWrapper;
 import com.ning.billing.util.entity.dao.EntitySqlDaoWrapperFactory;
 
 import com.google.common.base.Function;
-import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
 
 public class DefaultBlockingStateDao 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>() {
         @Override
         public int compare(final BlockingStateModelDao o1, final BlockingStateModelDao o2) {
@@ -75,20 +60,11 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
         }
     });
 
-    private static final Logger log = LoggerFactory.getLogger(DefaultBlockingStateDao.class);
-
-    // Lame to rely on the API at this (low) level, but we need information from subscription to insert events not on disk
-    private final SubscriptionBaseInternalApi subscriptionInternalApi;
     private final EntitySqlDaoTransactionalJdbiWrapper transactionalSqlDao;
     private final Clock clock;
-    private final EntitlementUtils entitlementUtils;
 
-    @Inject
-    public DefaultBlockingStateDao(final SubscriptionBaseInternalApi subscriptionBaseInternalApi, final IDBI dbi, final Clock clock,
-                                   final CacheControllerDispatcher cacheControllerDispatcher, final NonEntityDao nonEntityDao,
-                                   final EntitlementUtils entitlementUtils) {
-        this.subscriptionInternalApi = subscriptionBaseInternalApi;
-        this.entitlementUtils = entitlementUtils;
+    public DefaultBlockingStateDao(final IDBI dbi, final Clock clock,
+                                   final CacheControllerDispatcher cacheControllerDispatcher, final NonEntityDao nonEntityDao) {
         this.transactionalSqlDao = new EntitySqlDaoTransactionalJdbiWrapper(dbi, clock, cacheControllerDispatcher, nonEntityDao);
         this.clock = clock;
     }
@@ -131,14 +107,13 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
             @Override
             public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
                 final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final List<BlockingStateModelDao> modelsOnDisk = sqlDao.getBlockingHistoryForService(blockableId, serviceName, context);
-                final List<BlockingStateModelDao> allModels = addBlockingStatesNotOnDisk(blockableId, modelsOnDisk, sqlDao, context);
-                return new ArrayList<BlockingState>(Collections2.transform(allModels, new Function<BlockingStateModelDao, BlockingState>() {
-                    @Override
-                    public BlockingState apply(@Nullable final BlockingStateModelDao src) {
-                        return BlockingStateModelDao.toBlockingState(src);
-                    }
-                }));
+                return new ArrayList<BlockingState>(Collections2.transform(sqlDao.getBlockingHistoryForService(blockableId, serviceName, context),
+                                                                           new Function<BlockingStateModelDao, BlockingState>() {
+                                                                               @Override
+                                                                               public BlockingState apply(@Nullable final BlockingStateModelDao src) {
+                                                                                   return BlockingStateModelDao.toBlockingState(src);
+                                                                               }
+                                                                           }));
             }
         });
     }
@@ -149,14 +124,13 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
             @Override
             public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
                 final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final List<BlockingStateModelDao> modelsOnDisk = sqlDao.getBlockingAll(blockableId, context);
-                final List<BlockingStateModelDao> allModels = addBlockingStatesNotOnDisk(blockableId, modelsOnDisk, sqlDao, context);
-                return new ArrayList<BlockingState>(Collections2.transform(allModels, new Function<BlockingStateModelDao, BlockingState>() {
-                    @Override
-                    public BlockingState apply(@Nullable final BlockingStateModelDao src) {
-                        return BlockingStateModelDao.toBlockingState(src);
-                    }
-                }));
+                return new ArrayList<BlockingState>(Collections2.transform(sqlDao.getBlockingAll(blockableId, context),
+                                                                           new Function<BlockingStateModelDao, BlockingState>() {
+                                                                               @Override
+                                                                               public BlockingState apply(@Nullable final BlockingStateModelDao src) {
+                                                                                   return BlockingStateModelDao.toBlockingState(src);
+                                                                               }
+                                                                           }));
             }
         });
     }
@@ -167,14 +141,13 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
             @Override
             public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
                 final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final List<BlockingStateModelDao> modelsOnDisk = sqlDao.getByAccountRecordId(context);
-                final List<BlockingStateModelDao> allModels = addBlockingStatesNotOnDisk(null, modelsOnDisk, sqlDao, context);
-                return new ArrayList<BlockingState>(Collections2.transform(allModels, new Function<BlockingStateModelDao, BlockingState>() {
-                    @Override
-                    public BlockingState apply(@Nullable final BlockingStateModelDao src) {
-                        return BlockingStateModelDao.toBlockingState(src);
-                    }
-                }));
+                return new ArrayList<BlockingState>(Collections2.transform(sqlDao.getByAccountRecordId(context),
+                                                                           new Function<BlockingStateModelDao, BlockingState>() {
+                                                                               @Override
+                                                                               public BlockingState apply(@Nullable final BlockingStateModelDao src) {
+                                                                                   return BlockingStateModelDao.toBlockingState(src);
+                                                                               }
+                                                                           }));
             }
         });
     }
@@ -240,74 +213,4 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
             }
         });
     }
-
-    // Add blocking states for add-ons, which would be impacted by a future cancellation or change of their base plan
-    // See DefaultEntitlement#blockAddOnsIfRequired
-    private List<BlockingStateModelDao> addBlockingStatesNotOnDisk(@Nullable final UUID blockableId,
-                                                                   final List<BlockingStateModelDao> blockingStatesOnDisk,
-                                                                   final BlockingStateSqlDao sqlDao,
-                                                                   final InternalTenantContext context) {
-        final Collection<BlockingStateModelDao> blockingStatesOnDiskCopy = new LinkedList<BlockingStateModelDao>(blockingStatesOnDisk);
-
-        // Find all base entitlements that we care about (for which we want to find future cancelled add-ons)
-        final Iterable<SubscriptionBase> baseSubscriptionsToConsider;
-        try {
-            if (blockableId == null) {
-                // We're coming from getBlockingAllForAccountRecordId
-                final Iterable<SubscriptionBase> subscriptions = Iterables.<SubscriptionBase>concat(subscriptionInternalApi.getSubscriptionsForAccount(context).values());
-                baseSubscriptionsToConsider = Iterables.<SubscriptionBase>filter(subscriptions,
-                                                                                 new Predicate<SubscriptionBase>() {
-                                                                                     @Override
-                                                                                     public boolean apply(final SubscriptionBase input) {
-                                                                                         return ProductCategory.BASE.equals(input.getCategory()) &&
-                                                                                                !EntitlementState.CANCELLED.equals(input.getState());
-                                                                                     }
-                                                                                 });
-            } else {
-                // We're coming from getBlockingHistoryForService / getBlockingAll, but we don't know the blocking type
-                final SubscriptionBase addOnSubscription;
-                try {
-                    addOnSubscription = subscriptionInternalApi.getSubscriptionFromId(blockableId, context);
-                } catch (SubscriptionBaseApiException ignored) {
-                    // blockable id points to an account or bundle, in which case there are no extra blocking states to add
-                    return blockingStatesOnDisk;
-                }
-
-                // blockable id points to a subscription, but make sure it's an add-on
-                if (ProductCategory.ADD_ON.equals(addOnSubscription.getCategory())) {
-                    final SubscriptionBase baseSubscription = subscriptionInternalApi.getBaseSubscription(addOnSubscription.getBundleId(), context);
-                    baseSubscriptionsToConsider = ImmutableList.<SubscriptionBase>of(baseSubscription);
-                } else {
-                    // blockable id points to a base or standalone subscription, there is nothing to do
-                    return blockingStatesOnDisk;
-                }
-            }
-        } catch (SubscriptionBaseApiException e) {
-            log.error("Error retrieving subscriptions for account record id " + context.getAccountRecordId(), e);
-            throw new RuntimeException(e);
-        }
-
-        final DateTime now = clock.getUTCNow();
-        for (final SubscriptionBase subscriptionBase : baseSubscriptionsToConsider) {
-            final Collection<BlockingState> blockingStates;
-            try {
-                blockingStates = entitlementUtils.computeFutureBlockingStatesForAssociatedAddons(sqlDao.getBlockingHistoryForService(subscriptionBase.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, context),
-                                                                                                 subscriptionBase,
-                                                                                                 now,
-                                                                                                 context);
-            } catch (EntitlementApiException e) {
-                log.error("Error computing blocking states for addons for account record id " + context.getAccountRecordId(), e);
-                throw new RuntimeException(e);
-            }
-
-            // Inject the extra blocking states into the stream
-            for (final BlockingState blockingState : blockingStates) {
-                final BlockingStateModelDao blockingStateModelDao = new BlockingStateModelDao(blockingState, now, now);
-                blockingStatesOnDiskCopy.add(blockingStateModelDao);
-            }
-        }
-
-        // Return the sorted list
-        return BLOCKING_STATE_MODEL_DAO_ORDERING.immutableSortedCopy(blockingStatesOnDiskCopy);
-    }
 }
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
new file mode 100644
index 0000000..a35b1f3
--- /dev/null
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
@@ -0,0 +1,199 @@
+/*
+ * Copyright 2010-2013 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.entitlement.dao;
+
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.UUID;
+
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+import javax.inject.Singleton;
+
+import org.joda.time.DateTime;
+import org.skife.jdbi.v2.IDBI;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.ning.billing.callcontext.InternalCallContext;
+import com.ning.billing.callcontext.InternalTenantContext;
+import com.ning.billing.catalog.api.ProductCategory;
+import com.ning.billing.clock.Clock;
+import com.ning.billing.entitlement.api.BlockingState;
+import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
+import com.ning.billing.entitlement.api.EntitlementApiException;
+import com.ning.billing.entitlement.engine.core.EventsStream;
+import com.ning.billing.entitlement.engine.core.EventsStreamBuilder;
+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.util.cache.CacheControllerDispatcher;
+import com.ning.billing.util.dao.NonEntityDao;
+
+import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
+
+@Singleton
+public class ProxyBlockingStateDao implements BlockingStateDao {
+
+    private static final Logger log = LoggerFactory.getLogger(ProxyBlockingStateDao.class);
+
+    private static final Ordering<BlockingState> BLOCKING_STATE_ORDERING = Ordering.<BlockingState>from(new Comparator<BlockingState>() {
+        @Override
+        public int compare(final BlockingState o1, final BlockingState o2) {
+            final int blockableIdComparison = o1.getBlockedId().compareTo(o2.getBlockedId());
+            if (blockableIdComparison == 0) {
+                // effective_date column NOT NULL
+                final int comparison = o1.getEffectiveDate().compareTo(o2.getEffectiveDate());
+                if (comparison == 0) {
+                    // Keep a stable ordering for ties
+                    return o1.getCreatedDate().compareTo(o2.getCreatedDate());
+                } else {
+                    return comparison;
+                }
+            } else {
+                return blockableIdComparison;
+            }
+        }
+    });
+
+    private final EventsStreamBuilder eventsStreamBuilder;
+    private final SubscriptionBaseInternalApi subscriptionInternalApi;
+    private final Clock clock;
+    private final DefaultBlockingStateDao delegate;
+
+    @Inject
+    public ProxyBlockingStateDao(final EventsStreamBuilder eventsStreamBuilder, final SubscriptionBaseInternalApi subscriptionBaseInternalApi,
+                                 final IDBI dbi, final Clock clock,
+                                 final CacheControllerDispatcher cacheControllerDispatcher, final NonEntityDao nonEntityDao) {
+        this.eventsStreamBuilder = eventsStreamBuilder;
+        this.subscriptionInternalApi = subscriptionBaseInternalApi;
+        this.clock = clock;
+        this.delegate = new DefaultBlockingStateDao(dbi, clock, cacheControllerDispatcher, nonEntityDao);
+    }
+
+    @Override
+    public BlockingState getBlockingStateForService(final UUID blockableId, final String serviceName, final InternalTenantContext context) {
+        return delegate.getBlockingStateForService(blockableId, serviceName, context);
+    }
+
+    @Override
+    public List<BlockingState> getBlockingState(final UUID blockableId, final InternalTenantContext context) {
+        return delegate.getBlockingState(blockableId, context);
+    }
+
+    @Override
+    public List<BlockingState> getBlockingHistoryForService(final UUID blockableId, final String serviceName, final InternalTenantContext context) {
+        final List<BlockingState> statesOnDisk = delegate.getBlockingHistoryForService(blockableId, serviceName, context);
+        return addBlockingStatesNotOnDisk(blockableId, statesOnDisk, context);
+    }
+
+    @Override
+    public List<BlockingState> getBlockingAll(final UUID blockableId, final InternalTenantContext context) {
+        final List<BlockingState> statesOnDisk = delegate.getBlockingAll(blockableId, context);
+        return addBlockingStatesNotOnDisk(blockableId, statesOnDisk, context);
+    }
+
+    @Override
+    public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
+        final List<BlockingState> statesOnDisk = delegate.getBlockingAllForAccountRecordId(context);
+        return addBlockingStatesNotOnDisk(null, statesOnDisk, context);
+    }
+
+    @Override
+    public void setBlockingState(final BlockingState state, final Clock clock, final InternalCallContext context) {
+        delegate.setBlockingState(state, clock, context);
+    }
+
+    @Override
+    public void unactiveBlockingState(final UUID blockableId, final InternalCallContext context) {
+        delegate.unactiveBlockingState(blockableId, context);
+    }
+
+    // Add blocking states for add-ons, which would be impacted by a future cancellation or change of their base plan
+    // See DefaultEntitlement#blockAddOnsIfRequired
+    private List<BlockingState> addBlockingStatesNotOnDisk(@Nullable final UUID blockableId,
+                                                           final List<BlockingState> blockingStatesOnDisk,
+                                                           final InternalTenantContext context) {
+        final Collection<BlockingState> blockingStatesOnDiskCopy = new LinkedList<BlockingState>(blockingStatesOnDisk);
+
+        // Find all base entitlements that we care about (for which we want to find future cancelled add-ons)
+        final Iterable<SubscriptionBase> baseSubscriptionsToConsider;
+        try {
+            if (blockableId == null) {
+                // We're coming from getBlockingAllForAccountRecordId
+                final Iterable<SubscriptionBase> subscriptions = Iterables.<SubscriptionBase>concat(subscriptionInternalApi.getSubscriptionsForAccount(context).values());
+                baseSubscriptionsToConsider = Iterables.<SubscriptionBase>filter(subscriptions,
+                                                                                 new Predicate<SubscriptionBase>() {
+                                                                                     @Override
+                                                                                     public boolean apply(final SubscriptionBase input) {
+                                                                                         return ProductCategory.BASE.equals(input.getCategory()) &&
+                                                                                                !EntitlementState.CANCELLED.equals(input.getState());
+                                                                                     }
+                                                                                 });
+            } else {
+                // We're coming from getBlockingHistoryForService / getBlockingAll, but we don't know the blocking type
+                final SubscriptionBase addOnSubscription;
+                try {
+                    addOnSubscription = subscriptionInternalApi.getSubscriptionFromId(blockableId, context);
+                } catch (SubscriptionBaseApiException ignored) {
+                    // blockable id points to an account or bundle, in which case there are no extra blocking states to add
+                    return blockingStatesOnDisk;
+                }
+
+                // blockable id points to a subscription, but make sure it's an add-on
+                if (ProductCategory.ADD_ON.equals(addOnSubscription.getCategory())) {
+                    final SubscriptionBase baseSubscription = subscriptionInternalApi.getBaseSubscription(addOnSubscription.getBundleId(), context);
+                    baseSubscriptionsToConsider = ImmutableList.<SubscriptionBase>of(baseSubscription);
+                } else {
+                    // blockable id points to a base or standalone subscription, there is nothing to do
+                    return blockingStatesOnDisk;
+                }
+            }
+        } catch (SubscriptionBaseApiException e) {
+            log.error("Error retrieving subscriptions for account record id " + context.getAccountRecordId(), e);
+            throw new RuntimeException(e);
+        }
+
+        // Compute the blocking states not on disk for all base subscriptions
+        final DateTime now = clock.getUTCNow();
+        for (final SubscriptionBase baseSubscription : baseSubscriptionsToConsider) {
+            final EventsStream eventsStream;
+            try {
+                eventsStream = eventsStreamBuilder.buildForEntitlement(baseSubscription.getId(), context);
+            } catch (EntitlementApiException e) {
+                log.error("Error computing blocking states for addons for account record id " + context.getAccountRecordId(), e);
+                throw new RuntimeException(e);
+            }
+
+            final Collection<BlockingState> blockingStatesNotOnDisk = eventsStream.computeAddonsBlockingStatesForFutureSubscriptionBaseEvents();
+
+            // Inject the extra blocking states into the stream
+            for (final BlockingState blockingState : blockingStatesNotOnDisk) {
+                final BlockingStateModelDao blockingStateModelDao = new BlockingStateModelDao(blockingState, now, now);
+                blockingStatesOnDiskCopy.add(BlockingStateModelDao.toBlockingState(blockingStateModelDao));
+            }
+        }
+
+        // Return the sorted list
+        return BLOCKING_STATE_ORDERING.immutableSortedCopy(blockingStatesOnDiskCopy);
+    }
+}
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 22b4f3b..9c440da 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
@@ -17,11 +17,8 @@
 package com.ning.billing.entitlement.engine.core;
 
 import java.io.IOException;
-import java.util.Collection;
-import java.util.List;
 import java.util.UUID;
 
-import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.joda.time.DateTime;
@@ -32,46 +29,24 @@ 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.callcontext.InternalTenantContext;
-import com.ning.billing.catalog.api.Plan;
-import com.ning.billing.catalog.api.Product;
-import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.DefaultEntitlementService;
-import com.ning.billing.entitlement.EntitlementService;
 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.api.DefaultBlockingTransitionInternalEvent;
-import com.ning.billing.entitlement.api.DefaultEntitlementApi;
-import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
-import com.ning.billing.entitlement.api.EntitlementApiException;
 import com.ning.billing.entitlement.block.BlockingChecker;
 import com.ning.billing.entitlement.block.BlockingChecker.BlockingAggregator;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
-import com.ning.billing.entitlement.dao.BlockingStateModelDao;
-import com.ning.billing.junction.DefaultBlockingState;
 import com.ning.billing.notificationq.api.NotificationEvent;
 import com.ning.billing.notificationq.api.NotificationQueue;
 import com.ning.billing.notificationq.api.NotificationQueueService;
 import com.ning.billing.notificationq.api.NotificationQueueService.NoSuchNotificationQueue;
-import com.ning.billing.subscription.api.SubscriptionBase;
-import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
-import com.ning.billing.subscription.api.SubscriptionBaseTransitionType;
-import com.ning.billing.subscription.api.user.SubscriptionBaseTransition;
-
-import com.google.common.base.Function;
-import com.google.common.base.Predicate;
-import com.google.common.collect.Collections2;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
 
 public class EntitlementUtils {
 
     private static final Logger log = LoggerFactory.getLogger(EntitlementUtils.class);
 
-    private final SubscriptionBaseInternalApi subscriptionInternalApi;
     private final BlockingStateDao dao;
     private final BlockingChecker blockingChecker;
     private final PersistentBus eventBus;
@@ -79,11 +54,9 @@ public class EntitlementUtils {
     protected final NotificationQueueService notificationQueueService;
 
     @Inject
-    public EntitlementUtils(final SubscriptionBaseInternalApi subscriptionInternalApi,
-                            final BlockingStateDao dao, final BlockingChecker blockingChecker,
+    public EntitlementUtils(final BlockingStateDao dao, final BlockingChecker blockingChecker,
                             final PersistentBus eventBus, final Clock clock,
                             final NotificationQueueService notificationQueueService) {
-        this.subscriptionInternalApi = subscriptionInternalApi;
         this.dao = dao;
         this.blockingChecker = blockingChecker;
         this.eventBus = eventBus;
@@ -92,133 +65,6 @@ public class EntitlementUtils {
     }
 
     /**
-     * Compute future blocking states for addons associated to a base subscription following a change or a cancellation
-     * on that base subscription
-     * <p/>
-     * This is only used in the "read" path, to add events that are not on disk.
-     * See same logic in DefaultSubscriptionDao#buildBundleSubscriptions
-     *
-     * @param blockingStates        existing entitlement blocking states for that base subscription
-     * @param subscriptionBase      base subscription, reflecting the latest state (cancellation or change applied)
-     * @param now                   present reference time (to avoid timing issues in DefaultEntitlement#blockAddOnsIfRequired)
-     * @param internalTenantContext context  @return the blocking states for the add-ons
-     * @throws EntitlementApiException
-     */
-    public Collection<BlockingState> computeFutureBlockingStatesForAssociatedAddons(final Iterable<BlockingStateModelDao> blockingStates,
-                                                                                    final SubscriptionBase subscriptionBase,
-                                                                                    final DateTime now,
-                                                                                    final InternalTenantContext internalTenantContext) throws EntitlementApiException {
-        if (!ProductCategory.BASE.equals(subscriptionBase.getCategory())) {
-            // Only base subscriptions have add-ons
-            return ImmutableList.<BlockingState>of();
-        }
-
-        // We need to find the first "trigger" transition, from which we will create the add-ons cancellation events.
-        // This can either be a future entitlement cancel...
-        final BlockingStateModelDao futureEntitlementCancelEvent = Iterables.<BlockingStateModelDao>tryFind(blockingStates,
-                                                                                                            new Predicate<BlockingStateModelDao>() {
-                                                                                                                @Override
-                                                                                                                public boolean apply(final BlockingStateModelDao input) {
-                                                                                                                    // Look at future cancellations only
-                                                                                                                    return input.getEffectiveDate().isAfter(now) &&
-                                                                                                                           EntitlementService.ENTITLEMENT_SERVICE_NAME.equals(input.getService()) &&
-                                                                                                                           DefaultEntitlementApi.ENT_STATE_CANCELLED.equals(input.getState());
-                                                                                                                }
-                                                                                                            }).orNull();
-
-        if (futureEntitlementCancelEvent != null) {
-            // Note that in theory we could always only look subscription base as we assume entitlement cancel means subscription base cancel
-            // but we want to use the effective date of the entitlement cancel event to create the add-on cancel event
-            return computeBlockingStatesForAssociatedAddons(subscriptionBase.getBundleId(), null, futureEntitlementCancelEvent.getEffectiveDate(), internalTenantContext);
-        } else {
-            // ...or a subscription change (i.e. a change plan where the new plan has an impact on the existing add-on).
-            // We need to go back to subscription base as entitlement doesn't know about these
-            final SubscriptionBaseTransition futureSubscriptionBaseChangeEvent = Iterables.<SubscriptionBaseTransition>tryFind(subscriptionBase.getAllTransitions(),
-                                                                                                                               new Predicate<SubscriptionBaseTransition>() {
-                                                                                                                                   @Override
-                                                                                                                                   public boolean apply(final SubscriptionBaseTransition input) {
-                                                                                                                                       // Look at future changes only
-                                                                                                                                       return input.getEffectiveTransitionTime().isAfter(now) &&
-                                                                                                                                              (SubscriptionBaseTransitionType.CHANGE.equals(input.getTransitionType()) ||
-                                                                                                                                               // This should never happen, as we should always have an entitlement cancel event
-                                                                                                                                               // (see above), but add it just in case...
-                                                                                                                                               SubscriptionBaseTransitionType.CANCEL.equals(input.getTransitionType()));
-                                                                                                                                   }
-                                                                                                                               }).orNull();
-            if (futureSubscriptionBaseChangeEvent == null) {
-                // Nothing to do
-                return ImmutableList.<BlockingState>of();
-            }
-
-            final Plan nextPlan = futureSubscriptionBaseChangeEvent.getNextPlan();
-            final Product product = nextPlan == null ? futureSubscriptionBaseChangeEvent.getPreviousPlan().getProduct() : nextPlan.getProduct();
-            return computeBlockingStatesForAssociatedAddons(subscriptionBase.getBundleId(), product, futureSubscriptionBaseChangeEvent.getEffectiveTransitionTime(), internalTenantContext);
-        }
-    }
-
-    // "write" path (when the cancellation/change is effective)
-    public Collection<BlockingState> computeBlockingStatesForAssociatedAddons(final SubscriptionBase subscriptionBase, final DateTime effectiveDate, final InternalTenantContext internalTenantContext) {
-        return computeBlockingStatesForAssociatedAddons(subscriptionBase.getBundleId(), EntitlementState.CANCELLED.equals(subscriptionBase.getState()) ? null : subscriptionBase.getLastActiveProduct(), effectiveDate, internalTenantContext);
-    }
-
-    private Collection<BlockingState> computeBlockingStatesForAssociatedAddons(final UUID bundleId,
-                                                                               @Nullable final Product baseTransitionTriggerNextProduct,
-                                                                               final DateTime baseTransitionTriggerEffectiveTransitionTime,
-                                                                               final InternalTenantContext internalTenantContext) {
-        final Collection<String> includedAddonsForProduct;
-        final Collection<String> availableAddonsForProduct;
-        if (baseTransitionTriggerNextProduct == null) {
-            includedAddonsForProduct = ImmutableList.<String>of();
-            availableAddonsForProduct = ImmutableList.<String>of();
-        } else {
-            includedAddonsForProduct = Collections2.<Product, String>transform(ImmutableSet.<Product>copyOf(baseTransitionTriggerNextProduct.getIncluded()),
-                                                                               new Function<Product, String>() {
-                                                                                   @Override
-                                                                                   public String apply(final Product product) {
-                                                                                       return product.getName();
-                                                                                   }
-                                                                               });
-
-            availableAddonsForProduct = Collections2.<Product, String>transform(ImmutableSet.<Product>copyOf(baseTransitionTriggerNextProduct.getAvailable()),
-                                                                                new Function<Product, String>() {
-                                                                                    @Override
-                                                                                    public String apply(final Product product) {
-                                                                                        return product.getName();
-                                                                                    }
-                                                                                });
-        }
-
-        // Retrieve all add-ons to block for that base subscription
-        final List<SubscriptionBase> subscriptionsForBundle = subscriptionInternalApi.getSubscriptionsForBundle(bundleId, internalTenantContext);
-        final Collection<SubscriptionBase> futureBlockedAddons = Collections2.<SubscriptionBase>filter(subscriptionsForBundle,
-                                                                                                       new Predicate<SubscriptionBase>() {
-                                                                                                           @Override
-                                                                                                           public boolean apply(final SubscriptionBase subscription) {
-                                                                                                               return ProductCategory.ADD_ON.equals(subscription.getCategory()) &&
-                                                                                                                      (
-                                                                                                                              // Base entitlement cancelled, cancel all add-ons
-                                                                                                                              // We don't check if the associated entitlement had already blocked here
-                                                                                                                              // but the dao should eventually do the right thing (it won't insert duplicated
-                                                                                                                              // blocking events)
-                                                                                                                              baseTransitionTriggerNextProduct == null ||
-                                                                                                                              (
-                                                                                                                                      // Change plan - check which add-ons to cancel
-                                                                                                                                      includedAddonsForProduct.contains(subscription.getLastActivePlan().getProduct().getName()) ||
-                                                                                                                                      !availableAddonsForProduct.contains(subscription.getLastActivePlan().getProduct().getName())
-                                                                                                                              )
-                                                                                                                      );
-                                                                                                           }
-                                                                                                       });
-        return Collections2.<SubscriptionBase, BlockingState>transform(futureBlockedAddons,
-                                                                       new Function<SubscriptionBase, BlockingState>() {
-                                                                           @Override
-                                                                           public BlockingState apply(final SubscriptionBase input) {
-                                                                               return new DefaultBlockingState(input.getId(), BlockingStateType.SUBSCRIPTION, DefaultEntitlementApi.ENT_STATE_CANCELLED, EntitlementService.ENTITLEMENT_SERVICE_NAME, true, true, false, baseTransitionTriggerEffectiveTransitionTime);
-                                                                           }
-                                                                       });
-    }
-
-    /**
      * Wrapper around BlockingStateDao#setBlockingState which will send an event on the bus if needed
      *
      * @param state   new state to store
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStream.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStream.java
index 99551db..9d07620 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStream.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStream.java
@@ -16,25 +16,34 @@
 
 package com.ning.billing.entitlement.engine.core;
 
+import java.util.Collection;
 import java.util.List;
+import java.util.UUID;
 
 import org.joda.time.DateTime;
 import org.joda.time.LocalDate;
 
 import com.ning.billing.account.api.Account;
 import com.ning.billing.callcontext.InternalTenantContext;
+import com.ning.billing.catalog.api.Product;
+import com.ning.billing.catalog.api.ProductCategory;
+import com.ning.billing.entitlement.EntitlementService;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
 import com.ning.billing.entitlement.api.DefaultEntitlementApi;
 import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 import com.ning.billing.entitlement.block.BlockingChecker.BlockingAggregator;
+import com.ning.billing.junction.DefaultBlockingState;
 import com.ning.billing.subscription.api.SubscriptionBase;
 import com.ning.billing.subscription.api.SubscriptionBaseTransitionType;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.subscription.api.user.SubscriptionBaseTransition;
 
+import com.google.common.base.Function;
 import com.google.common.base.Predicate;
+import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 
 public class EventsStream {
@@ -45,8 +54,12 @@ public class EventsStream {
     private final List<BlockingState> bundleEntitlementStates;
     private final List<BlockingState> accountEntitlementStates;
     private final BlockingAggregator blockingAggregator;
+    // Base subscription for the bundle if it exists, null otherwise
     private final SubscriptionBase baseSubscription;
+    // Subscription associated with this entitlement (equals to baseSubscription for base subscriptions)
     private final SubscriptionBase subscription;
+    // All subscriptions for that bundle
+    private final List<SubscriptionBase> allSubscriptionsForBundle;
     private final InternalTenantContext internalTenantContext;
     private final DateTime utcNow;
 
@@ -58,7 +71,7 @@ public class EventsStream {
                         final List<BlockingState> subscriptionEntitlementStates, final List<BlockingState> bundleEntitlementStates,
                         final List<BlockingState> accountEntitlementStates, final BlockingAggregator blockingAggregator,
                         final SubscriptionBase baseSubscription, final SubscriptionBase subscription,
-                        final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
+                        final List<SubscriptionBase> allSubscriptionsForBundle, final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
         this.account = account;
         this.bundle = bundle;
         this.subscriptionEntitlementStates = subscriptionEntitlementStates;
@@ -67,6 +80,7 @@ public class EventsStream {
         this.blockingAggregator = blockingAggregator;
         this.baseSubscription = baseSubscription;
         this.subscription = subscription;
+        this.allSubscriptionsForBundle = allSubscriptionsForBundle;
         this.internalTenantContext = contextWithValidAccountRecordId;
         this.utcNow = utcNow;
 
@@ -105,7 +119,7 @@ public class EventsStream {
         return blockingAggregator;
     }
 
-    public boolean isFutureEntitlementCancelled() {
+    public boolean isEntitlementFutureCancelled() {
         return entitlementCancelEvent != null && entitlementCancelEvent.getEffectiveDate().isAfter(utcNow);
     }
 
@@ -131,7 +145,7 @@ public class EventsStream {
                                                new Predicate<BlockingState>() {
                                                    @Override
                                                    public boolean apply(final BlockingState input) {
-                                                       return input.getEffectiveDate().isAfter(utcNow) &&
+                                                       return !input.getEffectiveDate().isBefore(utcNow) &&
                                                               typeList.contains(input.getStateName()) &&
                                                               (
                                                                       // ... for that subscription
@@ -147,18 +161,134 @@ public class EventsStream {
                                                });
     }
 
+    public BlockingState getEntitlementCancellationEvent(final UUID subscriptionId) {
+        return Iterables.<BlockingState>tryFind(subscriptionEntitlementStates,
+                                                new Predicate<BlockingState>() {
+                                                    @Override
+                                                    public boolean apply(final BlockingState input) {
+                                                        return DefaultEntitlementApi.ENT_STATE_CANCELLED.equals(input.getStateName()) &&
+                                                               input.getBlockedId().equals(subscriptionId);
+                                                    }
+                                                }).orNull();
+    }
+
     public Iterable<SubscriptionBaseTransition> getPendingSubscriptionEvents(final SubscriptionBaseTransitionType... types) {
+        return getPendingSubscriptionEvents(utcNow, types);
+    }
+
+    public Iterable<SubscriptionBaseTransition> getPendingSubscriptionEvents(final DateTime effectiveDatetime, final SubscriptionBaseTransitionType... types) {
         final List<SubscriptionBaseTransitionType> typeList = ImmutableList.<SubscriptionBaseTransitionType>copyOf(types);
         return Iterables.<SubscriptionBaseTransition>filter(subscription.getAllTransitions(),
                                                             new Predicate<SubscriptionBaseTransition>() {
                                                                 @Override
                                                                 public boolean apply(final SubscriptionBaseTransition input) {
-                                                                    return input.getEffectiveTransitionTime().isAfter(utcNow) &&
+                                                                    // Make sure we return the event for equality
+                                                                    return !input.getEffectiveTransitionTime().isBefore(effectiveDatetime) &&
                                                                            typeList.contains(input.getTransitionType());
                                                                 }
                                                             });
     }
 
+    public Collection<BlockingState> computeAddonsBlockingStatesForNextSubscriptionBaseEvent(final DateTime effectiveDate) {
+        // Compute the transition trigger
+        final Iterable<SubscriptionBaseTransition> pendingSubscriptionBaseTransitions = getPendingSubscriptionEvents(effectiveDate, SubscriptionBaseTransitionType.CHANGE, SubscriptionBaseTransitionType.CANCEL);
+        if (!pendingSubscriptionBaseTransitions.iterator().hasNext()) {
+            return ImmutableList.<BlockingState>of();
+        }
+
+        final SubscriptionBaseTransition subscriptionBaseTransitionTrigger = pendingSubscriptionBaseTransitions.iterator().next();
+        return computeAddonsBlockingStatesForSubscriptionBaseEvent(subscriptionBaseTransitionTrigger);
+    }
+
+    // Compute future blocking states not on disk for add-ons associated to this (base) events stream
+    public Collection<BlockingState> computeAddonsBlockingStatesForFutureSubscriptionBaseEvents() {
+        if (!ProductCategory.BASE.equals(subscription.getCategory())) {
+            // Only base subscriptions have add-ons
+            return ImmutableList.of();
+        }
+
+        // We need to find the first "trigger" transition, from which we will create the add-ons cancellation events.
+        // This can either be a future entitlement cancel...
+        if (isEntitlementFutureCancelled()) {
+            // Note that in theory we could always only look subscription base as we assume entitlement cancel means subscription base cancel
+            // but we want to use the effective date of the entitlement cancel event to create the add-on cancel event
+            final BlockingState futureEntitlementCancelEvent = getEntitlementCancellationEvent(subscription.getId());
+            return computeAddonsBlockingStatesForNextSubscriptionBaseEvent(futureEntitlementCancelEvent.getEffectiveDate());
+        } else {
+            // ...or a subscription change (i.e. a change plan where the new plan has an impact on the existing add-on).
+            // We need to go back to subscription base as entitlement doesn't know about these
+            return computeAddonsBlockingStatesForNextSubscriptionBaseEvent(utcNow);
+        }
+    }
+
+    private Collection<BlockingState> computeAddonsBlockingStatesForSubscriptionBaseEvent(final SubscriptionBaseTransition subscriptionBaseTransitionTrigger) {
+        if (baseSubscription == null || baseSubscription.getLastActivePlan() == null || !ProductCategory.BASE.equals(baseSubscription.getLastActivePlan().getProduct().getCategory())) {
+            return ImmutableList.<BlockingState>of();
+        }
+
+        final Product baseTransitionTriggerNextProduct = EntitlementState.CANCELLED.equals(subscriptionBaseTransitionTrigger.getNextState()) ? null : subscriptionBaseTransitionTrigger.getNextPlan().getProduct();
+
+        // Compute included and available addons for the new product
+        final Collection<String> includedAddonsForProduct;
+        final Collection<String> availableAddonsForProduct;
+        if (baseTransitionTriggerNextProduct == null) {
+            includedAddonsForProduct = ImmutableList.<String>of();
+            availableAddonsForProduct = ImmutableList.<String>of();
+        } else {
+            includedAddonsForProduct = Collections2.<Product, String>transform(ImmutableSet.<Product>copyOf(baseTransitionTriggerNextProduct.getIncluded()),
+                                                                               new Function<Product, String>() {
+                                                                                   @Override
+                                                                                   public String apply(final Product product) {
+                                                                                       return product.getName();
+                                                                                   }
+                                                                               });
+
+            availableAddonsForProduct = Collections2.<Product, String>transform(ImmutableSet.<Product>copyOf(baseTransitionTriggerNextProduct.getAvailable()),
+                                                                                new Function<Product, String>() {
+                                                                                    @Override
+                                                                                    public String apply(final Product product) {
+                                                                                        return product.getName();
+                                                                                    }
+                                                                                });
+        }
+
+        // Retrieve all add-ons to block for that base subscription
+        final Collection<SubscriptionBase> futureBlockedAddons = Collections2.<SubscriptionBase>filter(allSubscriptionsForBundle,
+                                                                                                       new Predicate<SubscriptionBase>() {
+                                                                                                           @Override
+                                                                                                           public boolean apply(final SubscriptionBase subscription) {
+                                                                                                               return ProductCategory.ADD_ON.equals(subscription.getCategory()) &&
+                                                                                                                      // Check the entitlement for that add-on hasn't been cancelled yet
+                                                                                                                      getEntitlementCancellationEvent(subscription.getId()) == null &&
+                                                                                                                      (
+                                                                                                                              // Base subscription cancelled
+                                                                                                                              baseTransitionTriggerNextProduct == null ||
+                                                                                                                              (
+                                                                                                                                      // Change plan - check which add-ons to cancel
+                                                                                                                                      includedAddonsForProduct.contains(subscription.getLastActivePlan().getProduct().getName()) ||
+                                                                                                                                      !availableAddonsForProduct.contains(subscription.getLastActivePlan().getProduct().getName())
+                                                                                                                              )
+                                                                                                                      );
+                                                                                                           }
+                                                                                                       });
+
+        // Create the blocking states
+        return Collections2.<SubscriptionBase, BlockingState>transform(futureBlockedAddons,
+                                                                       new Function<SubscriptionBase, BlockingState>() {
+                                                                           @Override
+                                                                           public BlockingState apply(final SubscriptionBase input) {
+                                                                               return new DefaultBlockingState(input.getId(),
+                                                                                                               BlockingStateType.SUBSCRIPTION,
+                                                                                                               DefaultEntitlementApi.ENT_STATE_CANCELLED,
+                                                                                                               EntitlementService.ENTITLEMENT_SERVICE_NAME,
+                                                                                                               true,
+                                                                                                               true,
+                                                                                                               false,
+                                                                                                               subscriptionBaseTransitionTrigger.getEffectiveTransitionTime());
+                                                                           }
+                                                                       });
+    }
+
     private void setup() {
         computeEntitlementEffectiveEndDate();
         computeEntitlementCancelEvent();
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
index 55785c5..a9d8578 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
@@ -22,10 +22,12 @@ import java.util.UUID;
 import javax.inject.Inject;
 import javax.inject.Singleton;
 
+import com.ning.billing.ObjectType;
 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.catalog.api.ProductCategory;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.EntitlementService;
 import com.ning.billing.entitlement.api.BlockingApiException;
@@ -41,6 +43,9 @@ import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.TenantContext;
 
+import com.google.common.base.Predicate;
+import com.google.common.collect.Iterables;
+
 @Singleton
 public class EventsStreamBuilder {
 
@@ -68,52 +73,64 @@ public class EventsStreamBuilder {
     }
 
     public EventsStream buildForBaseSubscription(final UUID bundleId, final TenantContext tenantContext) throws EntitlementApiException {
-        final SubscriptionBaseBundle bundle;
-        final SubscriptionBase subscription;
+        final SubscriptionBase baseSubscription;
         try {
             final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
-            bundle = subscriptionInternalApi.getBundleFromId(bundleId, internalTenantContext);
-            subscription = subscriptionInternalApi.getBaseSubscription(bundleId, internalTenantContext);
+            baseSubscription = subscriptionInternalApi.getBaseSubscription(bundleId, internalTenantContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        return buildForEntitlement(bundle, subscription, subscription, tenantContext);
+        return buildForEntitlement(baseSubscription.getId(), tenantContext);
     }
 
     public EventsStream buildForEntitlement(final UUID entitlementId, final TenantContext tenantContext) throws EntitlementApiException {
-        final SubscriptionBase baseSubscription;
-        final SubscriptionBase subscription;
+        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(entitlementId, ObjectType.SUBSCRIPTION, tenantContext);
+        return buildForEntitlement(entitlementId, internalTenantContext);
+    }
+
+    public EventsStream buildForEntitlement(final UUID entitlementId, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         final SubscriptionBaseBundle bundle;
+        final SubscriptionBase subscription;
+        final List<SubscriptionBase> allSubscriptionsForBundle;
+        final SubscriptionBase baseSubscription;
         try {
-            final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
             subscription = subscriptionInternalApi.getSubscriptionFromId(entitlementId, internalTenantContext);
             bundle = subscriptionInternalApi.getBundleFromId(subscription.getBundleId(), internalTenantContext);
-            baseSubscription = subscriptionInternalApi.getBaseSubscription(bundle.getId(), internalTenantContext);
+            allSubscriptionsForBundle = subscriptionInternalApi.getSubscriptionsForBundle(subscription.getBundleId(), internalTenantContext);
+            baseSubscription = Iterables.<SubscriptionBase>tryFind(allSubscriptionsForBundle,
+                                                                   new Predicate<SubscriptionBase>() {
+                                                                       @Override
+                                                                       public boolean apply(final SubscriptionBase input) {
+                                                                           return ProductCategory.BASE.equals(input.getLastActiveProduct().getCategory());
+                                                                       }
+                                                                   }).orNull(); // null for standalone subscriptions
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        return buildForEntitlement(bundle, baseSubscription, subscription, tenantContext);
+        return buildForEntitlement(bundle, baseSubscription, subscription, allSubscriptionsForBundle, internalTenantContext);
     }
 
-    private EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase baseSubscription, final SubscriptionBase subscription, final TenantContext tenantContext) throws EntitlementApiException {
-        final InternalTenantContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalTenantContext(bundle.getAccountId(), tenantContext);
-
+    private EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle,
+                                             final SubscriptionBase baseSubscription,
+                                             final SubscriptionBase subscription,
+                                             final List<SubscriptionBase> allSubscriptionsForBundle,
+                                             final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         final Account account;
         try {
-            account = accountInternalApi.getAccountById(bundle.getAccountId(), contextWithValidAccountRecordId);
+            account = accountInternalApi.getAccountById(bundle.getAccountId(), internalTenantContext);
         } catch (AccountApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        final List<BlockingState> subscriptionEntitlementStates = blockingStateDao.getBlockingHistoryForService(subscription.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, contextWithValidAccountRecordId);
-        final List<BlockingState> bundleEntitlementStates = blockingStateDao.getBlockingHistoryForService(bundle.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, contextWithValidAccountRecordId);
-        final List<BlockingState> accountEntitlementStates = blockingStateDao.getBlockingHistoryForService(account.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, contextWithValidAccountRecordId);
+        final List<BlockingState> subscriptionEntitlementStates = blockingStateDao.getBlockingHistoryForService(subscription.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext);
+        final List<BlockingState> bundleEntitlementStates = blockingStateDao.getBlockingHistoryForService(bundle.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext);
+        final List<BlockingState> accountEntitlementStates = blockingStateDao.getBlockingHistoryForService(account.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext);
 
         final BlockingAggregator blockingAggregator;
         try {
-            blockingAggregator = checker.getBlockedStatus(subscription, contextWithValidAccountRecordId);
+            blockingAggregator = checker.getBlockedStatus(subscription, internalTenantContext);
         } catch (BlockingApiException e) {
             throw new EntitlementApiException(e);
         }
@@ -126,7 +143,8 @@ public class EventsStreamBuilder {
                                 blockingAggregator,
                                 baseSubscription,
                                 subscription,
-                                contextWithValidAccountRecordId,
+                                allSubscriptionsForBundle,
+                                internalTenantContext,
                                 clock.getUTCNow());
     }
 }
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/glue/DefaultEntitlementModule.java b/entitlement/src/main/java/com/ning/billing/entitlement/glue/DefaultEntitlementModule.java
index f639db3..d95248b 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/glue/DefaultEntitlementModule.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/glue/DefaultEntitlementModule.java
@@ -28,7 +28,7 @@ import com.ning.billing.entitlement.api.svcs.DefaultInternalBlockingApi;
 import com.ning.billing.entitlement.block.BlockingChecker;
 import com.ning.billing.entitlement.block.DefaultBlockingChecker;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
-import com.ning.billing.entitlement.dao.DefaultBlockingStateDao;
+import com.ning.billing.entitlement.dao.ProxyBlockingStateDao;
 import com.ning.billing.entitlement.engine.core.EntitlementUtils;
 import com.ning.billing.entitlement.engine.core.EventsStreamBuilder;
 import com.ning.billing.glue.EntitlementModule;
@@ -56,7 +56,7 @@ public class DefaultEntitlementModule extends AbstractModule implements Entitlem
 
     @Override
     public void installBlockingStateDao() {
-        bind(BlockingStateDao.class).to(DefaultBlockingStateDao.class).asEagerSingleton();
+        bind(BlockingStateDao.class).to(ProxyBlockingStateDao.class).asEagerSingleton();
     }
 
     @Override
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 feb040c..0f7c8e0 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
@@ -104,7 +104,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, addOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime);
         // and for the "write" path (which will be exercised when the future notification kicks in).
         // Note that no event are computed because the add-on is not cancelled yet
-        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, addOnEntitlement, null);
+        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, addOnEntitlement, null, false);
         // Verify also the blocking states DAO adds events not on disk
         checkBlockingStatesDAO(baseEntitlement, addOnEntitlement, baseEffectiveCancellationOrChangeDate, true);
 
@@ -121,14 +121,14 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, null);
         // ...and for the "write" path (which has been exercised when the notification kicked in).
-        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime);
+        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime, false);
         // Verify also the blocking states API doesn't add too many events (now on disk)
         checkBlockingStatesDAO(cancelledBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveCancellationOrChangeDate, true);
     }
 
     @Test(groups = "slow", description = "Verify add-ons blocking states are not impacted by IMM cancellations")
     public void testCancellationIMM() throws Exception {
-        // A bit fragile? The blocking state check (checkBlockingStatesDAO) could be a bit off
+        // Approximate check, as the blocking state check (checkBlockingStatesDAO) could be a bit off
         final DateTime cancellationDateTime = clock.getUTCNow();
         final LocalDate cancellationDate = clock.getUTCToday();
 
@@ -145,7 +145,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, null);
         // ...and for the "write" path (which has been exercised in the cancel call above).
-        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDateTime);
+        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDateTime, true);
         // Verify also the blocking states DAO doesn't add too many events (all on disk)
         checkBlockingStatesDAO(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDate, true);
 
@@ -156,7 +156,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, null);
-        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDateTime);
+        checkActualBlockingStatesToCancel(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDateTime, true);
         checkBlockingStatesDAO(cancelledBaseEntitlement, cancelledAddOnEntitlement, cancellationDate, true);
     }
 
@@ -172,7 +172,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, addOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime);
         // ...and for the "write" path (which will be exercised when the future notification kicks in).
         // Note that no event are computed because the add-on is not cancelled yet
-        checkActualBlockingStatesToCancel(changedBaseEntitlement, addOnEntitlement, null);
+        checkActualBlockingStatesToCancel(changedBaseEntitlement, addOnEntitlement, null, false);
         // Verify also the blocking states DAO adds events not on disk
         checkBlockingStatesDAO(changedBaseEntitlement, addOnEntitlement, baseEffectiveCancellationOrChangeDate, false);
 
@@ -189,14 +189,14 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, null);
         // ...and for the "write" path (which has been exercised when the notification kicked in).
-        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime);
+        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveEOTCancellationOrChangeDateTime, false);
         // Verify also the blocking states API doesn't add too many events (now on disk)
         checkBlockingStatesDAO(changedBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveCancellationOrChangeDate, false);
     }
 
     @Test(groups = "slow", description = "Verify add-ons blocking states are added for IMM change plans")
     public void testChangePlanIMM() throws Exception {
-        // A bit fragile? The blocking state check (checkBlockingStatesDAO) could be a bit off
+        // Approximate check, as the blocking state check (checkBlockingStatesDAO) could be a bit off
         final DateTime changeDateTime = clock.getUTCNow();
         final LocalDate changeDate = clock.getUTCToday();
 
@@ -213,7 +213,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, null);
         // ...and for the "write" path (which has been exercised in the change call above).
-        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, changeDateTime);
+        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, changeDateTime, true);
         // Verify also the blocking states DAO doesn't add too many events (all on disk)
         checkBlockingStatesDAO(changedBaseEntitlement, cancelledAddOnEntitlement, changeDate, false);
 
@@ -224,7 +224,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, null, null);
         checkFutureBlockingStatesToCancel(cancelledAddOnEntitlement, null, null);
         checkFutureBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, null);
-        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, changeDateTime);
+        checkActualBlockingStatesToCancel(changedBaseEntitlement, cancelledAddOnEntitlement, changeDateTime, true);
         checkBlockingStatesDAO(changedBaseEntitlement, cancelledAddOnEntitlement, changeDate, false);
     }
 
@@ -245,7 +245,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
     }
 
     // Test the "write" path
-    private void checkActualBlockingStatesToCancel(final DefaultEntitlement baseEntitlement, final DefaultEntitlement addOnEntitlement, @Nullable final DateTime effectiveCancellationDateTime) throws EntitlementApiException {
+    private void checkActualBlockingStatesToCancel(final DefaultEntitlement baseEntitlement, final DefaultEntitlement addOnEntitlement, @Nullable final DateTime effectiveCancellationDateTime, final boolean approximateDateCheck) throws EntitlementApiException {
         final Collection<BlockingState> blockingStatesForCancellation = computeBlockingStatesForAssociatedAddons(baseEntitlement, effectiveCancellationDateTime);
         if (effectiveCancellationDateTime == null) {
             Assert.assertEquals(blockingStatesForCancellation.size(), 0);
@@ -253,7 +253,12 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
             Assert.assertEquals(blockingStatesForCancellation.size(), 1);
             final BlockingState blockingState = blockingStatesForCancellation.iterator().next();
             Assert.assertEquals(blockingState.getBlockedId(), addOnEntitlement.getId());
-            Assert.assertEquals(blockingState.getEffectiveDate(), effectiveCancellationDateTime);
+            if (approximateDateCheck) {
+                Assert.assertEquals(blockingState.getEffectiveDate().toLocalDate(), effectiveCancellationDateTime.toLocalDate());
+                Assert.assertEquals(blockingState.getEffectiveDate().getMinuteOfDay(), effectiveCancellationDateTime.getMinuteOfDay());
+            } else {
+                Assert.assertEquals(blockingState.getEffectiveDate(), effectiveCancellationDateTime);
+            }
             Assert.assertEquals(blockingState.getType(), BlockingStateType.SUBSCRIPTION);
             Assert.assertEquals(blockingState.getService(), EntitlementService.ENTITLEMENT_SERVICE_NAME);
             Assert.assertEquals(blockingState.getStateName(), DefaultEntitlementApi.ENT_STATE_CANCELLED);
@@ -282,13 +287,12 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
     }
 
     private Collection<BlockingState> computeFutureBlockingStatesForAssociatedAddons(final DefaultEntitlement baseEntitlement) throws EntitlementApiException {
-        return entitlementUtils.computeFutureBlockingStatesForAssociatedAddons(sqlDao.getBlockingHistoryForService(baseEntitlement.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, internalCallContext),
-                                                                               baseEntitlement.getSubscriptionBase(),
-                                                                               clock.getUTCNow(),
-                                                                               internalCallContext);
+        final EventsStream eventsStream = eventsStreamBuilder.buildForEntitlement(baseEntitlement.getId(), callContext);
+        return eventsStream.computeAddonsBlockingStatesForFutureSubscriptionBaseEvents();
     }
 
     private Collection<BlockingState> computeBlockingStatesForAssociatedAddons(final DefaultEntitlement baseEntitlement, final DateTime effectiveDate) throws EntitlementApiException {
-        return entitlementUtils.computeBlockingStatesForAssociatedAddons(baseEntitlement.getSubscriptionBase(), effectiveDate, internalCallContext);
+        final EventsStream eventsStream = eventsStreamBuilder.buildForEntitlement(baseEntitlement.getId(), callContext);
+        return eventsStream.computeAddonsBlockingStatesForNextSubscriptionBaseEvent(effectiveDate);
     }
 }
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 b8b5768..9f7c6ec 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/EntitlementTestSuiteWithEmbeddedDB.java
@@ -44,6 +44,7 @@ import com.ning.billing.entitlement.api.EntitlementApi;
 import com.ning.billing.entitlement.api.SubscriptionApi;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.entitlement.engine.core.EntitlementUtils;
+import com.ning.billing.entitlement.engine.core.EventsStreamBuilder;
 import com.ning.billing.entitlement.glue.TestEntitlementModuleWithEmbeddedDB;
 import com.ning.billing.junction.BlockingInternalApi;
 import com.ning.billing.mock.MockAccountBuilder;
@@ -102,6 +103,8 @@ public class EntitlementTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuiteWi
     protected EntitlementService entitlementService;
     @Inject
     protected EntitlementUtils entitlementUtils;
+    @Inject
+    protected EventsStreamBuilder eventsStreamBuilder;
 
     protected Catalog catalog;
 
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseApiService.java b/subscription/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseApiService.java
index b6f5a7b..f718fef 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseApiService.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseApiService.java
@@ -52,17 +52,17 @@ public interface SubscriptionBaseApiService {
     public boolean uncancel(DefaultSubscriptionBase subscription, CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlan(DefaultSubscriptionBase subscription, String productName, BillingPeriod term,
                                String priceList, CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlanWithRequestedDate(DefaultSubscriptionBase subscription, String productName, BillingPeriod term,
                                                 String priceList, DateTime requestedDate, CallContext context)
             throws SubscriptionBaseApiException;
 
-    // Return the effective date of the change, null for immediate
+    // Return the effective date of the change
     public DateTime changePlanWithPolicy(DefaultSubscriptionBase subscription, String productName, BillingPeriod term,
                                          String priceList, BillingActionPolicy policy, CallContext context)
             throws SubscriptionBaseApiException;
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java b/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
index 87fb575..edb96ec 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
@@ -399,7 +399,7 @@ public class DefaultSubscriptionBaseApiService implements SubscriptionBaseApiSer
 
         final boolean isChangeImmediate = subscription.getCurrentPlan().getProduct().getName().equals(newProductName) &&
                                           subscription.getCurrentPlan().getBillingPeriod() == newBillingPeriod;
-        return isChangeImmediate ? null : effectiveDate;
+        return effectiveDate;
     }
 
 
diff --git a/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiError.java b/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiError.java
index 5761a86..bd50241 100644
--- a/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiError.java
+++ b/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiError.java
@@ -33,6 +33,7 @@ import com.ning.billing.catalog.api.BillingPeriod;
 import com.ning.billing.catalog.api.Duration;
 import com.ning.billing.catalog.api.PlanPhase;
 import com.ning.billing.catalog.api.PriceListSet;
+import com.ning.billing.clock.DefaultClock;
 import com.ning.billing.subscription.SubscriptionTestSuiteNoDB;
 import com.ning.billing.subscription.api.SubscriptionBase;
 import com.ning.billing.subscription.exceptions.SubscriptionBaseError;
@@ -157,7 +158,9 @@ public class TestUserApiError extends SubscriptionTestSuiteNoDB {
             assertEquals(subscriptionInternalApi.getSubscriptionFromId(subscription.getId(), internalCallContext).getCurrentPlan().getBillingPeriod(), BillingPeriod.ANNUAL);
         }
 
-        assertNull(subscription.changePlanWithPolicy("Shotgun", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, BillingActionPolicy.IMMEDIATE, callContext));
+        // Assume the call takes less than a second
+        assertEquals(DefaultClock.truncateMs(subscription.changePlanWithPolicy("Shotgun", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, BillingActionPolicy.IMMEDIATE, callContext)),
+                     DefaultClock.truncateMs(clock.getUTCNow()));
         assertEquals(subscriptionInternalApi.getSubscriptionFromId(subscription.getId(), internalCallContext).getCurrentPlan().getBillingPeriod(), BillingPeriod.MONTHLY);
     }