killbill-memoizeit

entitlement: refactor DefaultEntitlement Make it use the

11/8/2013 2:47:33 PM

Details

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 bcd3e0f..ddca8b5 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
@@ -18,19 +18,15 @@ package com.ning.billing.entitlement.api;
 
 import java.io.IOException;
 import java.util.Collection;
-import java.util.List;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
 
 import org.joda.time.DateTime;
-import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
 
 import com.ning.billing.ErrorCode;
 import com.ning.billing.account.api.Account;
-import com.ning.billing.account.api.AccountApiException;
-import com.ning.billing.account.api.AccountInternalApi;
 import com.ning.billing.callcontext.InternalCallContext;
 import com.ning.billing.catalog.api.BillingActionPolicy;
 import com.ning.billing.catalog.api.BillingPeriod;
@@ -48,6 +44,7 @@ import com.ning.billing.entitlement.engine.core.EntitlementNotificationKey;
 import com.ning.billing.entitlement.engine.core.EntitlementNotificationKeyAction;
 import com.ning.billing.entitlement.engine.core.EntitlementUtils;
 import com.ning.billing.entitlement.engine.core.EventsStream;
+import com.ning.billing.entitlement.engine.core.EventsStreamBuilder;
 import com.ning.billing.entity.EntityBase;
 import com.ning.billing.junction.DefaultBlockingState;
 import com.ning.billing.notificationq.api.NotificationEvent;
@@ -61,17 +58,13 @@ import com.ning.billing.util.callcontext.CallContext;
 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.Collections2;
-
 public class DefaultEntitlement extends EntityBase implements Entitlement {
 
+    protected final EventsStreamBuilder eventsStreamBuilder;
     protected final EntitlementDateHelper dateHelper;
     protected final InternalCallContextFactory internalCallContextFactory;
     protected final Clock clock;
     protected final BlockingChecker checker;
-    protected final UUID accountId;
-    protected final AccountInternalApi accountApi;
     protected final EntitlementApi entitlementApi;
     protected final SubscriptionBaseInternalApi subscriptionInternalApi;
     protected final BlockingStateDao blockingStateDao;
@@ -79,37 +72,30 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
     protected final EntitlementUtils entitlementUtils;
 
     // Refresh-able
-    protected SubscriptionBase subscriptionBase;
-    protected EntitlementState state;
-    protected LocalDate effectiveEndDate;
-    protected String externalKey;
-    protected DateTimeZone accountTimeZone;
-
-    public DefaultEntitlement(final EntitlementDateHelper dateHelper, final EventsStream eventsStream,
-                              final AccountInternalApi accountApi, final EntitlementApi entitlementApi, final SubscriptionBaseInternalApi subscriptionInternalApi, final InternalCallContextFactory internalCallContextFactory,
-                              final BlockingStateDao blockingStateDao, final Clock clock, final BlockingChecker checker, final NotificationQueueService notificationQueueService,
-                              final EntitlementUtils entitlementUtils) {
-        this(dateHelper, eventsStream.getSubscription(), eventsStream.getAccount().getId(), eventsStream.getBundle().getExternalKey(),
-             eventsStream.getEntitlementState(), eventsStream.getEntitlementEffectiveEndDate(), eventsStream.getAccount().getTimeZone(),
-             accountApi, entitlementApi, subscriptionInternalApi, internalCallContextFactory, blockingStateDao, clock, checker, notificationQueueService, entitlementUtils);
-    }
-
-    public DefaultEntitlement(final EntitlementDateHelper dateHelper, final SubscriptionBase subscriptionBase, final UUID accountId,
-                              final String externalKey, final EntitlementState state, final LocalDate effectiveEndDate, final DateTimeZone accountTimeZone,
-                              final AccountInternalApi accountApi, final EntitlementApi entitlementApi, final SubscriptionBaseInternalApi subscriptionInternalApi, final InternalCallContextFactory internalCallContextFactory,
-                              final BlockingStateDao blockingStateDao, final Clock clock, final BlockingChecker checker, final NotificationQueueService notificationQueueService,
-                              final EntitlementUtils entitlementUtils) {
-        super(subscriptionBase.getId(), subscriptionBase.getCreatedDate(), subscriptionBase.getUpdatedDate());
+    protected EventsStream eventsStream;
+
+    public DefaultEntitlement(final UUID entitlementId, final EventsStreamBuilder eventsStreamBuilder,
+                              final EntitlementApi entitlementApi, final BlockingStateDao blockingStateDao,
+                              final SubscriptionBaseInternalApi subscriptionInternalApi, final BlockingChecker checker,
+                              final NotificationQueueService notificationQueueService, final EntitlementUtils entitlementUtils,
+                              final EntitlementDateHelper dateHelper, final Clock clock,
+                              final InternalCallContextFactory internalCallContextFactory, final TenantContext tenantContext) throws EntitlementApiException {
+        this(eventsStreamBuilder.buildForEntitlement(entitlementId, tenantContext), eventsStreamBuilder,
+             entitlementApi, blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
+             entitlementUtils, dateHelper, clock, internalCallContextFactory);
+    }
+
+    public DefaultEntitlement(final EventsStream eventsStream, final EventsStreamBuilder eventsStreamBuilder,
+                              final EntitlementApi entitlementApi, final BlockingStateDao blockingStateDao,
+                              final SubscriptionBaseInternalApi subscriptionInternalApi, final BlockingChecker checker,
+                              final NotificationQueueService notificationQueueService, final EntitlementUtils entitlementUtils,
+                              final EntitlementDateHelper dateHelper, final Clock clock, final InternalCallContextFactory internalCallContextFactory) {
+        super(eventsStream.getSubscription().getId(), eventsStream.getSubscription().getCreatedDate(), eventsStream.getSubscription().getUpdatedDate());
+        this.eventsStreamBuilder = eventsStreamBuilder;
+        this.eventsStream = eventsStream;
         this.dateHelper = dateHelper;
-        this.subscriptionBase = subscriptionBase;
-        this.accountId = accountId;
-        this.externalKey = externalKey;
-        this.state = state;
-        this.effectiveEndDate = effectiveEndDate;
-        this.accountApi = accountApi;
         this.entitlementApi = entitlementApi;
         this.subscriptionInternalApi = subscriptionInternalApi;
-        this.accountTimeZone = accountTimeZone;
         this.internalCallContextFactory = internalCallContextFactory;
         this.clock = clock;
         this.checker = checker;
@@ -119,26 +105,33 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
     }
 
     public DefaultEntitlement(final DefaultEntitlement in) {
-        this(in.getDateHelper(),
-             in.getSubscriptionBase(),
-             in.getAccountId(),
-             in.getExternalKey(),
-             in.getState(),
-             in.getEffectiveEndDate(),
-             in.getAccountTimeZone(),
-             in.getAccountApi(),
+        this(in.getEventsStream(),
+             in.getEventsStreamBuilder(),
              in.getEntitlementApi(),
-             in.getSubscriptionInternalApi(),
-             in.getInternalCallContextFactory(),
              in.getBlockingStateDao(),
-             in.getClock(),
+             in.getSubscriptionInternalApi(),
              in.getChecker(),
              in.getNotificationQueueService(),
-             in.getEntitlementUtils());
+             in.getEntitlementUtils(),
+             in.getDateHelper(),
+             in.getClock(),
+             in.getInternalCallContextFactory());
+    }
+
+    public Account getAccount() {
+        return eventsStream.getAccount();
+    }
+
+    public EventsStream getEventsStream() {
+        return eventsStream;
     }
 
     public SubscriptionBase getSubscriptionBase() {
-        return subscriptionBase;
+        return eventsStream.getSubscription();
+    }
+
+    public EventsStreamBuilder getEventsStreamBuilder() {
+        return eventsStreamBuilder;
     }
 
     public EntitlementDateHelper getDateHelper() {
@@ -149,10 +142,6 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         return internalCallContextFactory;
     }
 
-    public AccountInternalApi getAccountApi() {
-        return accountApi;
-    }
-
     public EntitlementApi getEntitlementApi() {
         return entitlementApi;
     }
@@ -169,10 +158,6 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         return checker;
     }
 
-    public DateTimeZone getAccountTimeZone() {
-        return accountTimeZone;
-    }
-
     public BlockingStateDao getBlockingStateDao() {
         return blockingStateDao;
     }
@@ -187,76 +172,74 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public UUID getBaseEntitlementId() {
-        return subscriptionBase.getId();
+        return eventsStream.getSubscription().getId();
     }
 
     @Override
     public UUID getBundleId() {
-        return subscriptionBase.getBundleId();
+        return eventsStream.getSubscription().getBundleId();
     }
 
     @Override
     public UUID getAccountId() {
-        return accountId;
+        return eventsStream.getAccount().getId();
     }
 
     @Override
     public String getExternalKey() {
-        return externalKey;
+        return eventsStream.getBundle().getExternalKey();
     }
 
     @Override
     public EntitlementState getState() {
-        return state;
+        return eventsStream.getEntitlementState();
     }
 
     @Override
     public EntitlementSourceType getSourceType() {
-        return subscriptionBase.getSourceType();
+        return getSubscriptionBase().getSourceType();
     }
 
     @Override
     public LocalDate getEffectiveStartDate() {
-        return new LocalDate(subscriptionBase.getStartDate(), accountTimeZone);
+        return new LocalDate(getSubscriptionBase().getStartDate(), eventsStream.getAccount().getTimeZone());
     }
 
     @Override
     public LocalDate getEffectiveEndDate() {
-        return effectiveEndDate;
+        return eventsStream.getEntitlementEffectiveEndDate();
     }
 
     @Override
     public Product getLastActiveProduct() {
-        return subscriptionBase.getLastActiveProduct();
+        return getSubscriptionBase().getLastActiveProduct();
     }
 
     @Override
     public Plan getLastActivePlan() {
-        return subscriptionBase.getLastActivePlan();
+        return getSubscriptionBase().getLastActivePlan();
     }
 
     @Override
     public PlanPhase getLastActivePhase() {
-        return subscriptionBase.getLastActivePhase();
+        return getSubscriptionBase().getLastActivePhase();
     }
 
     @Override
     public PriceList getLastActivePriceList() {
-        return subscriptionBase.getLastActivePriceList();
+        return getSubscriptionBase().getLastActivePriceList();
     }
 
     @Override
     public ProductCategory getLastActiveProductCategory() {
-        return subscriptionBase.getLastActiveCategory();
+        return getSubscriptionBase().getLastActiveCategory();
     }
 
 
     @Override
     public Entitlement cancelEntitlementWithPolicy(final EntitlementActionPolicy entitlementPolicy, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk - required to have the latest CTD
-        refresh(callContext, contextWithValidAccountRecordId);
+        refresh(callContext);
 
         final LocalDate cancellationDate = getLocalDateFromEntitlementPolicy(entitlementPolicy);
         return cancelEntitlementWithDate(cancellationDate, false, callContext);
@@ -264,21 +247,20 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement cancelEntitlementWithDate(final LocalDate localCancelDate, final boolean overrideBillingEffectiveDate, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, contextWithValidAccountRecordId);
+        refresh(callContext);
 
-        if (state == EntitlementState.CANCELLED) {
+        if (eventsStream.isEntitlementCancelled()) {
             throw new EntitlementApiException(ErrorCode.SUB_CANCEL_BAD_STATE, getId(), EntitlementState.CANCELLED);
         }
 
-        final DateTime effectiveCancelDate = dateHelper.fromLocalDateAndReferenceTime(localCancelDate, subscriptionBase.getStartDate(), contextWithValidAccountRecordId);
+        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
+        final DateTime effectiveCancelDate = dateHelper.fromLocalDateAndReferenceTime(localCancelDate, getSubscriptionBase().getStartDate(), contextWithValidAccountRecordId);
         try {
             if (overrideBillingEffectiveDate) {
-                subscriptionBase.cancelWithDate(effectiveCancelDate, callContext);
+                getSubscriptionBase().cancelWithDate(effectiveCancelDate, callContext);
             } else {
-                subscriptionBase.cancel(callContext);
+                getSubscriptionBase().cancel(callContext);
             }
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
@@ -294,10 +276,8 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement cancelEntitlementWithPolicyOverrideBillingPolicy(final EntitlementActionPolicy entitlementPolicy, final BillingActionPolicy billingPolicy, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk - required to have the latest CTD
-        refresh(callContext, contextWithValidAccountRecordId);
+        refresh(callContext);
 
         final LocalDate cancellationDate = getLocalDateFromEntitlementPolicy(entitlementPolicy);
         return cancelEntitlementWithDateOverrideBillingPolicy(cancellationDate, billingPolicy, callContext);
@@ -305,46 +285,24 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public void uncancelEntitlement(final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, contextWithValidAccountRecordId);
+        refresh(callContext);
 
-        if (state == EntitlementState.CANCELLED || subscriptionBase.getState() == EntitlementState.CANCELLED) {
+        if (eventsStream.isEntitlementCancelled() || eventsStream.isSubscriptionCancelled()) {
             throw new EntitlementApiException(ErrorCode.SUB_CANCEL_BAD_STATE, getId(), EntitlementState.CANCELLED);
         }
-        final List<BlockingState> blockingStatesForAccount = blockingStateDao.getBlockingAllForAccountRecordId(contextWithValidAccountRecordId);
-        final Collection<BlockingState> futureEntitlementCancellationEvents = Collections2.filter(blockingStatesForAccount, new Predicate<BlockingState>() {
-            @Override
-            public boolean apply(final BlockingState input) {
-                // Delete all future cancellation events...
-                return EntitlementService.ENTITLEMENT_SERVICE_NAME.equals(input.getService()) &&
-                       DefaultEntitlementApi.ENT_STATE_CANCELLED.equals(input.getStateName()) &&
-                       input.getEffectiveDate().isAfter(clock.getUTCNow()) &&
-                       (
-                               // ... for that subscription
-                               BlockingStateType.SUBSCRIPTION.equals(input.getType()) && input.getBlockedId().equals(getId()) ||
-                               // ... for the associated base subscription (to make sure an add-on isn't associated with a cancelled base entitlement)
-                               BlockingStateType.SUBSCRIPTION.equals(input.getType()) && input.getBlockedId().equals(getBaseEntitlementId()) ||
-                               // ... for that bundle (to make sure the subscription isn't associated with a cancelled bundle - not yet implemented)
-                               BlockingStateType.SUBSCRIPTION_BUNDLE.equals(input.getType()) && input.getBlockedId().equals(getBundleId()) ||
-                               // ... for that bundle (to make sure the subscription isn't associated with a cancelled account - not yet implemented)
-                               BlockingStateType.ACCOUNT.equals(input.getType()) && input.getBlockedId().equals(getAccountId())
-                       );
-            }
-        });
 
-        // Reactivate entitlement
-        // We should only have one future event in theory - but cleanup the data if it's not the case
-        // See https://github.com/killbill/killbill/issues/111
-        for (final BlockingState futureCancellation : futureEntitlementCancellationEvents) {
+        // Reactivate entitlements
+        // See also https://github.com/killbill/killbill/issues/111
+        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
+        for (final BlockingState futureCancellation : eventsStream.getPendingEntitlementCancellationEvents()) {
             blockingStateDao.unactiveBlockingState(futureCancellation.getId(), contextWithValidAccountRecordId);
         }
 
         // If billing was previously cancelled, reactivate
-        if (subscriptionBase.getFutureEndDate() != null) {
+        if (getSubscriptionBase().getFutureEndDate() != null) {
             try {
-                subscriptionBase.uncancel(callContext);
+                getSubscriptionBase().uncancel(callContext);
             } catch (SubscriptionBaseApiException e) {
                 throw new EntitlementApiException(e);
             }
@@ -353,23 +311,22 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement cancelEntitlementWithDateOverrideBillingPolicy(final LocalDate localCancelDate, final BillingActionPolicy billingPolicy, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, contextWithValidAccountRecordId);
+        refresh(callContext);
 
-        if (state == EntitlementState.CANCELLED) {
+        if (eventsStream.isEntitlementCancelled()) {
             throw new EntitlementApiException(ErrorCode.SUB_CANCEL_BAD_STATE, getId(), EntitlementState.CANCELLED);
         }
 
         try {
-            subscriptionBase.cancelWithPolicy(billingPolicy, callContext);
+            getSubscriptionBase().cancelWithPolicy(billingPolicy, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        final LocalDate effectiveLocalDate = new LocalDate(localCancelDate, accountTimeZone);
-        final DateTime effectiveDate = dateHelper.fromLocalDateAndReferenceTime(effectiveLocalDate, subscriptionBase.getStartDate(), contextWithValidAccountRecordId);
+        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);
 
@@ -382,10 +339,10 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         final LocalDate cancellationDate;
         switch (entitlementPolicy) {
             case IMMEDIATE:
-                cancellationDate = new LocalDate(clock.getUTCNow(), accountTimeZone);
+                cancellationDate = new LocalDate(clock.getUTCNow(), eventsStream.getAccount().getTimeZone());
                 break;
             case END_OF_TERM:
-                cancellationDate = subscriptionBase.getChargedThroughDate() != null ? new LocalDate(subscriptionBase.getChargedThroughDate(), accountTimeZone) : new LocalDate(clock.getUTCNow(), accountTimeZone);
+                cancellationDate = getSubscriptionBase().getChargedThroughDate() != null ? new LocalDate(getSubscriptionBase().getChargedThroughDate(), eventsStream.getAccount().getTimeZone()) : new LocalDate(clock.getUTCNow(), eventsStream.getAccount().getTimeZone());
                 break;
             default:
                 throw new RuntimeException("Unsupported policy " + entitlementPolicy);
@@ -395,24 +352,23 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement changePlan(final String productName, final BillingPeriod billingPeriod, final String priceList, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, context);
+        refresh(callContext);
 
-        if (state != EntitlementState.ACTIVE) {
-            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), state);
+        if (!eventsStream.isEntitlementActive()) {
+            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), getState());
         }
 
+        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
         try {
-            checker.checkBlockedChange(subscriptionBase, context);
+            checker.checkBlockedChange(getSubscriptionBase(), context);
         } catch (BlockingApiException e) {
             throw new EntitlementApiException(e, e.getCode(), e.getMessage());
         }
 
         final DateTime effectiveChangeDate;
         try {
-            effectiveChangeDate = subscriptionBase.changePlan(productName, billingPeriod, priceList, callContext);
+            effectiveChangeDate = getSubscriptionBase().changePlan(productName, billingPeriod, priceList, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
@@ -424,24 +380,23 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement changePlanWithDate(final String productName, final BillingPeriod billingPeriod, final String priceList, final LocalDate localDate, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, context);
+        refresh(callContext);
 
-        if (state != EntitlementState.ACTIVE) {
-            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), state);
+        if (!eventsStream.isEntitlementActive()) {
+            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), getState());
         }
 
+        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
         try {
-            checker.checkBlockedChange(subscriptionBase, context);
+            checker.checkBlockedChange(getSubscriptionBase(), context);
         } catch (BlockingApiException e) {
             throw new EntitlementApiException(e, e.getCode(), e.getMessage());
         }
 
-        final DateTime effectiveChangeDate = dateHelper.fromLocalDateAndReferenceTime(localDate, subscriptionBase.getStartDate(), context);
+        final DateTime effectiveChangeDate = dateHelper.fromLocalDateAndReferenceTime(localDate, getSubscriptionBase().getStartDate(), context);
         try {
-            subscriptionBase.changePlanWithDate(productName, billingPeriod, priceList, effectiveChangeDate, callContext);
+            getSubscriptionBase().changePlanWithDate(productName, billingPeriod, priceList, effectiveChangeDate, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
@@ -453,24 +408,23 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
 
     @Override
     public Entitlement changePlanOverrideBillingPolicy(final String productName, final BillingPeriod billingPeriod, final String priceList, final LocalDate localDateX, final BillingActionPolicy actionPolicy, final CallContext callContext) throws EntitlementApiException {
-        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(accountId, callContext);
-
         // Get the latest state from disk
-        refresh(callContext, context);
+        refresh(callContext);
 
-        if (state != EntitlementState.ACTIVE) {
-            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), state);
+        if (!eventsStream.isEntitlementActive()) {
+            throw new EntitlementApiException(ErrorCode.SUB_CHANGE_NON_ACTIVE, getId(), getState());
         }
 
+        final InternalCallContext context = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
         try {
-            checker.checkBlockedChange(subscriptionBase, context);
+            checker.checkBlockedChange(getSubscriptionBase(), context);
         } catch (BlockingApiException e) {
             throw new EntitlementApiException(e, e.getCode(), e.getMessage());
         }
 
         final DateTime effectiveChangeDate;
         try {
-            effectiveChangeDate = subscriptionBase.changePlanWithPolicy(productName, billingPeriod, priceList, actionPolicy, callContext);
+            effectiveChangeDate = getSubscriptionBase().changePlanWithPolicy(productName, billingPeriod, priceList, actionPolicy, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
@@ -480,43 +434,26 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         return entitlementApi.getEntitlementForId(getId(), callContext);
     }
 
-    private void refresh(final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
-        try {
-            final Account account = accountApi.getAccountById(accountId, internalCallContext);
-            // Not really required today as we don't allow timezone changes, but do it nonetheless
-            // in case we change our mind later (and so we don't get surprises...).
-            accountTimeZone = account.getTimeZone();
-        } catch (AccountApiException e) {
-            throw new EntitlementApiException(e);
-        }
-
-        final Entitlement refreshedEntitlement = entitlementApi.getEntitlementForId(getId(), context);
-        state = refreshedEntitlement.getState();
-        effectiveEndDate = refreshedEntitlement.getEffectiveEndDate();
-        externalKey = refreshedEntitlement.getExternalKey();
-        if (refreshedEntitlement instanceof DefaultEntitlement) {
-            subscriptionBase = ((DefaultEntitlement) refreshedEntitlement).getSubscriptionBase();
-        } else {
-            throw new IllegalStateException("We only support DefaultEntitlement implementations");
-        }
+    private void refresh(final TenantContext context) throws EntitlementApiException {
+        eventsStream = eventsStreamBuilder.refresh(eventsStream, context);
     }
 
     public void blockAddOnsIfRequired(@Nullable final DateTime effectiveDateOrNull, final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
         // Optimization - bail early
-        if (!ProductCategory.BASE.equals(subscriptionBase.getCategory())) {
+        if (!ProductCategory.BASE.equals(getSubscriptionBase().getCategory())) {
             // Only base subscriptions have add-ons
             return;
         }
 
         // Get the latest state from disk (we just got cancelled or changed plan)
-        refresh(context, internalCallContext);
+        refresh(context);
 
         final DateTime now = clock.getUTCNow();
 
         // null means immediate
         final DateTime effectiveDate = effectiveDateOrNull == null ? now : effectiveDateOrNull;
 
-        final boolean isBaseEntitlementCancelled = EntitlementState.CANCELLED.equals(state);
+        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
@@ -531,7 +468,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
             return;
         }
 
-        final Collection<BlockingState> addOnsBlockingStates = entitlementUtils.computeBlockingStatesForAssociatedAddons(subscriptionBase, effectiveDate, internalCallContext);
+        final Collection<BlockingState> addOnsBlockingStates = entitlementUtils.computeBlockingStatesForAssociatedAddons(getSubscriptionBase(), effectiveDate, internalCallContext);
         for (final BlockingState addOnBlockingState : addOnsBlockingStates) {
             entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(addOnBlockingState, internalCallContext);
         }
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlementApi.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlementApi.java
index a39d224..35919a5 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlementApi.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultEntitlementApi.java
@@ -42,7 +42,6 @@ import com.ning.billing.catalog.api.PlanPhaseSpecifier;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.EntitlementService;
 import com.ning.billing.entitlement.EntitlementTransitionType;
-import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 import com.ning.billing.entitlement.block.BlockingChecker;
 import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.entitlement.engine.core.EntitlementUtils;
@@ -111,18 +110,17 @@ public class DefaultEntitlementApi implements EntitlementApi {
     public Entitlement createBaseEntitlement(final UUID accountId, final PlanPhaseSpecifier planPhaseSpecifier, final String externalKey, final LocalDate effectiveDate, final CallContext callContext) throws EntitlementApiException {
         final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(accountId, callContext);
         try {
-            final Account account = accountApi.getAccountById(accountId, contextWithValidAccountRecordId);
             final SubscriptionBaseBundle bundle = subscriptionInternalApi.createBundleForAccount(accountId, externalKey, contextWithValidAccountRecordId);
 
             final DateTime referenceTime = clock.getUTCNow();
             final DateTime requestedDate = dateHelper.fromLocalDateAndReferenceTime(effectiveDate, referenceTime, contextWithValidAccountRecordId);
             final SubscriptionBase subscription = subscriptionInternalApi.createSubscription(bundle.getId(), planPhaseSpecifier, requestedDate, contextWithValidAccountRecordId);
-            return new DefaultEntitlement(dateHelper, subscription, accountId, bundle.getExternalKey(), EntitlementState.ACTIVE, null, account.getTimeZone(), accountApi, this,
-                                          subscriptionInternalApi, internalCallContextFactory, blockingStateDao, clock, checker, notificationQueueService, entitlementUtils);
+
+            return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, this,
+                                          blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
+                                          entitlementUtils, dateHelper, clock, internalCallContextFactory, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
-        } catch (AccountApiException e) {
-            throw new EntitlementApiException(e);
         }
     }
 
@@ -146,8 +144,9 @@ public class DefaultEntitlementApi implements EntitlementApi {
             final InternalCallContext context = internalCallContextFactory.createInternalCallContext(callContext);
             final SubscriptionBase subscription = subscriptionInternalApi.createSubscription(bundleId, planPhaseSpecifier, requestedDate, context);
 
-            return new DefaultEntitlement(dateHelper, subscription, eventsStreamForBaseSubscription.getAccount().getId(), eventsStreamForBaseSubscription.getBundle().getExternalKey(), EntitlementState.ACTIVE, null, eventsStreamForBaseSubscription.getAccount().getTimeZone(),
-                                          accountApi, this, subscriptionInternalApi, internalCallContextFactory, blockingStateDao, clock, checker, notificationQueueService, entitlementUtils);
+            return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, this,
+                                          blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
+                                          entitlementUtils, dateHelper, clock, internalCallContextFactory, callContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
@@ -171,8 +170,9 @@ public class DefaultEntitlementApi implements EntitlementApi {
     @Override
     public Entitlement getEntitlementForId(final UUID uuid, final TenantContext tenantContext) throws EntitlementApiException {
         final EventsStream eventsStream = eventsStreamBuilder.buildForEntitlement(uuid, tenantContext);
-        return new DefaultEntitlement(dateHelper, eventsStream, accountApi, this, subscriptionInternalApi, internalCallContextFactory,
-                                      blockingStateDao, clock, checker, notificationQueueService, entitlementUtils);
+        return new DefaultEntitlement(eventsStream, eventsStreamBuilder, this,
+                                      blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
+                                      entitlementUtils, dateHelper, clock, internalCallContextFactory);
     }
 
     @Override
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscription.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscription.java
index 8c1dac4..8f49b8a 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscription.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscription.java
@@ -23,7 +23,6 @@ import org.joda.time.LocalDate;
 
 public class DefaultSubscription extends DefaultEntitlement implements Subscription {
 
-
     private final List<BlockingState> blockingStates;
 
     DefaultSubscription(final DefaultEntitlement entitlement, final List<BlockingState> blockingStates) {
@@ -33,18 +32,18 @@ public class DefaultSubscription extends DefaultEntitlement implements Subscript
 
     @Override
     public LocalDate getBillingStartDate() {
-        return new LocalDate(subscriptionBase.getStartDate(), accountTimeZone);
+        return new LocalDate(getSubscriptionBase().getStartDate(), getAccount().getTimeZone());
     }
 
     @Override
     public LocalDate getBillingEndDate() {
-        final DateTime futureOrCurrentEndDate = subscriptionBase.getEndDate() != null ? subscriptionBase.getEndDate() : subscriptionBase.getFutureEndDate();
-        return futureOrCurrentEndDate != null ? new LocalDate(futureOrCurrentEndDate, accountTimeZone) : null;
+        final DateTime futureOrCurrentEndDate = getSubscriptionBase().getEndDate() != null ? getSubscriptionBase().getEndDate() : getSubscriptionBase().getFutureEndDate();
+        return futureOrCurrentEndDate != null ? new LocalDate(futureOrCurrentEndDate, getAccount().getTimeZone()) : null;
     }
 
     @Override
     public LocalDate getChargedThroughDate() {
-        return subscriptionBase.getChargedThroughDate() != null ? new LocalDate(subscriptionBase.getChargedThroughDate(), accountTimeZone) : null;
+        return getSubscriptionBase().getChargedThroughDate() != null ? new LocalDate(getSubscriptionBase().getChargedThroughDate(), getAccount().getTimeZone()) : null;
     }
 
     @Override
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 acd9558..99551db 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
@@ -23,8 +23,8 @@ import org.joda.time.LocalDate;
 
 import com.ning.billing.account.api.Account;
 import com.ning.billing.callcontext.InternalTenantContext;
-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;
@@ -45,6 +45,7 @@ public class EventsStream {
     private final List<BlockingState> bundleEntitlementStates;
     private final List<BlockingState> accountEntitlementStates;
     private final BlockingAggregator blockingAggregator;
+    private final SubscriptionBase baseSubscription;
     private final SubscriptionBase subscription;
     private final InternalTenantContext internalTenantContext;
     private final DateTime utcNow;
@@ -56,13 +57,15 @@ public class EventsStream {
     public EventsStream(final Account account, final SubscriptionBaseBundle bundle,
                         final List<BlockingState> subscriptionEntitlementStates, final List<BlockingState> bundleEntitlementStates,
                         final List<BlockingState> accountEntitlementStates, final BlockingAggregator blockingAggregator,
-                        final SubscriptionBase subscription, final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
+                        final SubscriptionBase baseSubscription, final SubscriptionBase subscription,
+                        final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
         this.account = account;
         this.bundle = bundle;
         this.subscriptionEntitlementStates = subscriptionEntitlementStates;
         this.bundleEntitlementStates = bundleEntitlementStates;
         this.accountEntitlementStates = accountEntitlementStates;
         this.blockingAggregator = blockingAggregator;
+        this.baseSubscription = baseSubscription;
         this.subscription = subscription;
         this.internalTenantContext = contextWithValidAccountRecordId;
         this.utcNow = utcNow;
@@ -114,15 +117,46 @@ public class EventsStream {
         return entitlementState == EntitlementState.CANCELLED;
     }
 
-    public SubscriptionBaseTransition getPendingSubscriptionEvents(final SubscriptionBaseTransitionType... types) {
+    public boolean isSubscriptionCancelled() {
+        return subscription.getState() == EntitlementState.CANCELLED;
+    }
+
+    public Iterable<BlockingState> getPendingEntitlementCancellationEvents() {
+        return getPendingEntitlementEvents(DefaultEntitlementApi.ENT_STATE_CANCELLED);
+    }
+
+    public Iterable<BlockingState> getPendingEntitlementEvents(final String... types) {
+        final List<String> typeList = ImmutableList.<String>copyOf(types);
+        return Iterables.<BlockingState>filter(subscriptionEntitlementStates,
+                                               new Predicate<BlockingState>() {
+                                                   @Override
+                                                   public boolean apply(final BlockingState input) {
+                                                       return input.getEffectiveDate().isAfter(utcNow) &&
+                                                              typeList.contains(input.getStateName()) &&
+                                                              (
+                                                                      // ... for that subscription
+                                                                      BlockingStateType.SUBSCRIPTION.equals(input.getType()) && input.getBlockedId().equals(subscription.getId()) ||
+                                                                      // ... for the associated base subscription
+                                                                      BlockingStateType.SUBSCRIPTION.equals(input.getType()) && input.getBlockedId().equals(baseSubscription.getId()) ||
+                                                                      // ... for that bundle
+                                                                      BlockingStateType.SUBSCRIPTION_BUNDLE.equals(input.getType()) && input.getBlockedId().equals(bundle.getId()) ||
+                                                                      // ... for that account
+                                                                      BlockingStateType.ACCOUNT.equals(input.getType()) && input.getBlockedId().equals(account.getId())
+                                                              );
+                                                   }
+                                               });
+    }
+
+    public Iterable<SubscriptionBaseTransition> getPendingSubscriptionEvents(final SubscriptionBaseTransitionType... types) {
         final List<SubscriptionBaseTransitionType> typeList = ImmutableList.<SubscriptionBaseTransitionType>copyOf(types);
-        return Iterables.<SubscriptionBaseTransition>tryFind(subscription.getAllTransitions(),
-                                                             new Predicate<SubscriptionBaseTransition>() {
-                                                                 @Override
-                                                                 public boolean apply(final SubscriptionBaseTransition input) {
-                                                                     return input.getEffectiveTransitionTime().isAfter(utcNow) && typeList.contains(input.getTransitionType());
-                                                                 }
-                                                             }).orNull();
+        return Iterables.<SubscriptionBaseTransition>filter(subscription.getAllTransitions(),
+                                                            new Predicate<SubscriptionBaseTransition>() {
+                                                                @Override
+                                                                public boolean apply(final SubscriptionBaseTransition input) {
+                                                                    return input.getEffectiveTransitionTime().isAfter(utcNow) &&
+                                                                           typeList.contains(input.getTransitionType());
+                                                                }
+                                                            });
     }
 
     private void setup() {
@@ -160,8 +194,7 @@ public class EventsStream {
                                                                   new Predicate<BlockingState>() {
                                                                       @Override
                                                                       public boolean apply(final BlockingState input) {
-                                                                          return EntitlementService.ENTITLEMENT_SERVICE_NAME.equals(input.getService()) &&
-                                                                                 DefaultEntitlementApi.ENT_STATE_CANCELLED.equals(input.getStateName());
+                                                                          return DefaultEntitlementApi.ENT_STATE_CANCELLED.equals(input.getStateName());
                                                                       }
                                                                   }).orNull();
     }
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 4634b00..55785c5 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
@@ -63,6 +63,10 @@ public class EventsStreamBuilder {
         this.internalCallContextFactory = internalCallContextFactory;
     }
 
+    public EventsStream refresh(final EventsStream eventsStream, final TenantContext tenantContext) throws EntitlementApiException {
+        return buildForEntitlement(eventsStream.getSubscription().getId(), tenantContext);
+    }
+
     public EventsStream buildForBaseSubscription(final UUID bundleId, final TenantContext tenantContext) throws EntitlementApiException {
         final SubscriptionBaseBundle bundle;
         final SubscriptionBase subscription;
@@ -74,24 +78,26 @@ public class EventsStreamBuilder {
             throw new EntitlementApiException(e);
         }
 
-        return buildForEntitlement(bundle, subscription, tenantContext);
+        return buildForEntitlement(bundle, subscription, subscription, tenantContext);
     }
 
     public EventsStream buildForEntitlement(final UUID entitlementId, final TenantContext tenantContext) throws EntitlementApiException {
+        final SubscriptionBase baseSubscription;
         final SubscriptionBase subscription;
         final SubscriptionBaseBundle bundle;
         try {
             final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
             subscription = subscriptionInternalApi.getSubscriptionFromId(entitlementId, internalTenantContext);
             bundle = subscriptionInternalApi.getBundleFromId(subscription.getBundleId(), internalTenantContext);
+            baseSubscription = subscriptionInternalApi.getBaseSubscription(bundle.getId(), internalTenantContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        return buildForEntitlement(bundle, subscription, tenantContext);
+        return buildForEntitlement(bundle, baseSubscription, subscription, tenantContext);
     }
 
-    private EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase subscription, final TenantContext tenantContext) throws EntitlementApiException {
+    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);
 
         final Account account;
@@ -118,6 +124,7 @@ public class EventsStreamBuilder {
                                 bundleEntitlementStates,
                                 accountEntitlementStates,
                                 blockingAggregator,
+                                baseSubscription,
                                 subscription,
                                 contextWithValidAccountRecordId,
                                 clock.getUTCNow());