killbill-memoizeit

entitlement: fix regression in bundle timeline API After

12/2/2013 4:14:56 PM

Details

diff --git a/api/src/main/java/com/ning/billing/entitlement/EventsStream.java b/api/src/main/java/com/ning/billing/entitlement/EventsStream.java
index fc57b74..b570fd3 100644
--- a/api/src/main/java/com/ning/billing/entitlement/EventsStream.java
+++ b/api/src/main/java/com/ning/billing/entitlement/EventsStream.java
@@ -60,11 +60,8 @@ public interface EventsStream {
 
     Collection<BlockingState> getPendingEntitlementCancellationEvents();
 
-    Collection<BlockingState> getSubscriptionEntitlementStates();
-
-    Collection<BlockingState> getBundleEntitlementStates();
-
-    Collection<BlockingState> getAccountEntitlementStates();
+    // All blocking states for the account, associated bundle or subscription
+    Collection<BlockingState> getBlockingStates();
 
     Collection<BlockingState> computeAddonsBlockingStatesForNextSubscriptionBaseEvent(DateTime effectiveDate);
 
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
index d3b480e..253525e 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
@@ -17,6 +17,7 @@
 package com.ning.billing.entitlement.api;
 
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
@@ -46,9 +47,28 @@ import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
 
 public class DefaultSubscriptionApi implements SubscriptionApi {
 
+    private static final Comparator<SubscriptionBundle> SUBSCRIPTION_BUNDLE_COMPARATOR = new Comparator<SubscriptionBundle>() {
+        @Override
+        public int compare(final SubscriptionBundle o1, final SubscriptionBundle o2) {
+            final int compared = o1.getOriginalCreatedDate().compareTo(o2.getOriginalCreatedDate());
+            if (compared != 0) {
+                return compared;
+            } else {
+                final int compared2 = o1.getUpdatedDate().compareTo(o2.getUpdatedDate());
+                if (compared2 != 0) {
+                    return compared2;
+                } else {
+                    // Default, stable, ordering
+                    return o1.getId().compareTo(o2.getId());
+                }
+            }
+        }
+    };
+
     private final EntitlementInternalApi entitlementInternalApi;
     private final SubscriptionBaseInternalApi subscriptionInternalApi;
     private final InternalCallContextFactory internalCallContextFactory;
@@ -188,7 +208,8 @@ public class DefaultSubscriptionApi implements SubscriptionApi {
             bundles.add(subscriptionBundle);
         }
 
-        return bundles;
+        // Sort the results for predictability
+        return Ordering.<SubscriptionBundle>from(SUBSCRIPTION_BUNDLE_COMPARATOR).sortedCopy(bundles);
     }
 
     private Map<UUID, List<Subscription>> buildSubscriptionsFromEntitlements(final AccountEntitlements accountEntitlements) {
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionBundleTimeline.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionBundleTimeline.java
index 608e6b6..20148a2 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionBundleTimeline.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionBundleTimeline.java
@@ -73,9 +73,7 @@ public class DefaultSubscriptionBundleTimeline implements SubscriptionBundleTime
     public DefaultSubscriptionBundleTimeline(final DateTimeZone accountTimeZone, final UUID accountId, final UUID bundleId, final String externalKey, final Collection<Entitlement> entitlements) {
         final Collection<BlockingState> blockingStates = new HashSet<BlockingState>();
         for (final Entitlement entitlement : entitlements) {
-            blockingStates.addAll(((DefaultEntitlement) entitlement).getEventsStream().getSubscriptionEntitlementStates());
-            blockingStates.addAll(((DefaultEntitlement) entitlement).getEventsStream().getBundleEntitlementStates());
-            blockingStates.addAll(((DefaultEntitlement) entitlement).getEventsStream().getAccountEntitlementStates());
+            blockingStates.addAll(((DefaultEntitlement) entitlement).getEventsStream().getBlockingStates());
         }
         this.accountId = accountId;
         this.bundleId = bundleId;
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
index 604b488..d810635 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
@@ -48,12 +48,12 @@ public class OptimizedProxyBlockingStateDao extends ProxyBlockingStateDao {
     }
 
     // Special signature for EventsStreamBuilder to save some DAO calls
-    public List<BlockingState> getBlockingHistoryForService(final List<BlockingState> blockingStatesOnDisk,
-                                                            final SubscriptionBaseBundle bundle,
-                                                            @Nullable final SubscriptionBase baseSubscription,
-                                                            final SubscriptionBase subscription,
-                                                            final List<SubscriptionBase> allSubscriptionsForBundle,
-                                                            final InternalTenantContext context) throws EntitlementApiException {
+    public List<BlockingState> getBlockingHistory(final List<BlockingState> blockingStatesOnDisk,
+                                                  final SubscriptionBaseBundle bundle,
+                                                  @Nullable final SubscriptionBase baseSubscription,
+                                                  final SubscriptionBase subscription,
+                                                  final List<SubscriptionBase> allSubscriptionsForBundle,
+                                                  final InternalTenantContext context) throws EntitlementApiException {
         // blockable id points to a subscription, but make sure it's an add-on
         if (!ProductCategory.ADD_ON.equals(subscription.getCategory())) {
             // blockable id points to a base or standalone subscription, there is nothing to do
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/DefaultEventsStream.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/DefaultEventsStream.java
index 9b2f7fc..7a735cb 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/DefaultEventsStream.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/DefaultEventsStream.java
@@ -57,9 +57,8 @@ public class DefaultEventsStream implements EventsStream {
 
     private final Account account;
     private final SubscriptionBaseBundle bundle;
-    private final List<BlockingState> subscriptionEntitlementStates;
-    private final List<BlockingState> bundleEntitlementStates;
-    private final List<BlockingState> accountEntitlementStates;
+    // All blocking states for the account, associated bundle or subscription
+    private final List<BlockingState> blockingStates;
     private final BlockingChecker blockingChecker;
     // Base subscription for the bundle if it exists, null otherwise
     private final SubscriptionBase baseSubscription;
@@ -71,23 +70,23 @@ public class DefaultEventsStream implements EventsStream {
     private final DateTime utcNow;
 
     private BlockingAggregator blockingAggregator;
+    private List<BlockingState> subscriptionEntitlementStates;
+    private List<BlockingState> bundleEntitlementStates;
+    private List<BlockingState> accountEntitlementStates;
     private List<BlockingState> currentSubscriptionEntitlementBlockingStatesForServices;
     private List<BlockingState> currentBundleEntitlementBlockingStatesForServices;
-    private List<BlockingState> currentAccountEntitlementBlockingStateForServices;
+    private List<BlockingState> currentAccountEntitlementBlockingStatesForServices;
     private LocalDate entitlementEffectiveEndDate;
     private BlockingState entitlementCancelEvent;
     private EntitlementState entitlementState;
 
     public DefaultEventsStream(final Account account, final SubscriptionBaseBundle bundle,
-                               final List<BlockingState> subscriptionEntitlementStates, final List<BlockingState> bundleEntitlementStates,
-                               final List<BlockingState> accountEntitlementStates, final BlockingChecker blockingChecker,
+                               final List<BlockingState> blockingStates, final BlockingChecker blockingChecker,
                                @Nullable final SubscriptionBase baseSubscription, final SubscriptionBase subscription,
                                final List<SubscriptionBase> allSubscriptionsForBundle, final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
         this.account = account;
         this.bundle = bundle;
-        this.subscriptionEntitlementStates = subscriptionEntitlementStates;
-        this.bundleEntitlementStates = bundleEntitlementStates;
-        this.accountEntitlementStates = accountEntitlementStates;
+        this.blockingStates = blockingStates;
         this.blockingChecker = blockingChecker;
         this.baseSubscription = baseSubscription;
         this.subscription = subscription;
@@ -98,10 +97,6 @@ public class DefaultEventsStream implements EventsStream {
         setup();
     }
 
-    public Account getAccount() {
-        return account;
-    }
-
     @Override
     public DateTimeZone getAccountTimeZone() {
         return account.getTimeZone();
@@ -127,10 +122,6 @@ public class DefaultEventsStream implements EventsStream {
         return subscription.getId();
     }
 
-    public SubscriptionBaseBundle getBundle() {
-        return bundle;
-    }
-
     @Override
     public SubscriptionBase getBaseSubscription() {
         return baseSubscription;
@@ -151,19 +142,11 @@ public class DefaultEventsStream implements EventsStream {
         return entitlementEffectiveEndDate;
     }
 
-    public BlockingState getEntitlementCancelEvent() {
-        return entitlementCancelEvent;
-    }
-
     @Override
     public EntitlementState getEntitlementState() {
         return entitlementState;
     }
 
-    public BlockingAggregator getCurrentBlockingAggregator() {
-        return blockingAggregator;
-    }
-
     @Override
     public boolean isBlockChange() {
         return blockingAggregator.isBlockChange();
@@ -198,18 +181,8 @@ public class DefaultEventsStream implements EventsStream {
     }
 
     @Override
-    public Collection<BlockingState> getSubscriptionEntitlementStates() {
-        return subscriptionEntitlementStates;
-    }
-
-    @Override
-    public Collection<BlockingState> getBundleEntitlementStates() {
-        return bundleEntitlementStates;
-    }
-
-    @Override
-    public Collection<BlockingState> getAccountEntitlementStates() {
-        return accountEntitlementStates;
+    public Collection<BlockingState> getBlockingStates() {
+        return blockingStates;
     }
 
     @Override
@@ -250,10 +223,6 @@ public class DefaultEventsStream implements EventsStream {
                                                 }).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(),
@@ -390,6 +359,7 @@ public class DefaultEventsStream implements EventsStream {
     }
 
     private void setup() {
+        computeEntitlementBlockingStates();
         computeBlockingAggregator();
         computeEntitlementEffectiveEndDate();
         computeEntitlementCancelEvent();
@@ -397,9 +367,9 @@ public class DefaultEventsStream implements EventsStream {
     }
 
     private void computeBlockingAggregator() {
-        final List<BlockingState> currentAccountEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(accountEntitlementStates);
-        final List<BlockingState> currentBundleEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(bundleEntitlementStates);
-        final List<BlockingState> currentSubscriptionEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(subscriptionEntitlementStates);
+        currentAccountEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(accountEntitlementStates);
+        currentBundleEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(bundleEntitlementStates);
+        currentSubscriptionEntitlementBlockingStatesForServices = filterCurrentBlockableStatePerService(subscriptionEntitlementStates);
         blockingAggregator = blockingChecker.getBlockedStatus(currentAccountEntitlementBlockingStatesForServices,
                                                               currentBundleEntitlementBlockingStatesForServices,
                                                               currentSubscriptionEntitlementBlockingStatesForServices,
@@ -465,4 +435,22 @@ public class DefaultEventsStream implements EventsStream {
             entitlementState = (blockingAggregator != null && blockingAggregator.isBlockEntitlement() ? EntitlementState.BLOCKED : EntitlementState.ACTIVE);
         }
     }
+
+    private void computeEntitlementBlockingStates() {
+        subscriptionEntitlementStates = filterBlockingStatesForEntitlementService(BlockingStateType.SUBSCRIPTION, subscription.getId());
+        bundleEntitlementStates = filterBlockingStatesForEntitlementService(BlockingStateType.SUBSCRIPTION_BUNDLE, subscription.getBundleId());
+        accountEntitlementStates = filterBlockingStatesForEntitlementService(BlockingStateType.ACCOUNT, account.getId());
+    }
+
+    private List<BlockingState> filterBlockingStatesForEntitlementService(final BlockingStateType blockingStateType, @Nullable final UUID blockableId) {
+        return ImmutableList.<BlockingState>copyOf(Iterables.<BlockingState>filter(blockingStates,
+                                                                                   new Predicate<BlockingState>() {
+                                                                                       @Override
+                                                                                       public boolean apply(final BlockingState input) {
+                                                                                           return blockingStateType.equals(input.getType()) &&
+                                                                                                  EntitlementService.ENTITLEMENT_SERVICE_NAME.equals(input.getService()) &&
+                                                                                                  input.getBlockedId().equals(blockableId);
+                                                                                       }
+                                                                                   }));
+    }
 }
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 f9f9cfb..be8a899 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
@@ -18,6 +18,7 @@ package com.ning.billing.entitlement.engine.core;
 
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -37,7 +38,6 @@ import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.entitlement.AccountEventsStreams;
-import com.ning.billing.entitlement.EntitlementService;
 import com.ning.billing.entitlement.EventsStream;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
@@ -144,28 +144,26 @@ public class EventsStreamBuilder {
         }
 
         // Retrieve the blocking states
-        final List<BlockingState> blockingStatesForAccount = BLOCKING_STATE_ORDERING.immutableSortedCopy(defaultBlockingStateDao.getBlockingAllForAccountRecordId(internalTenantContext));
+        final List<BlockingState> blockingStatesForAccount = defaultBlockingStateDao.getBlockingAllForAccountRecordId(internalTenantContext);
 
-        // Optimization: build lookup tables for entitlement states
-        final List<BlockingState> accountEntitlementStates = new LinkedList<BlockingState>();
-        final Map<UUID, List<BlockingState>> entitlementStatesPerSubscription = new HashMap<UUID, List<BlockingState>>();
-        final Map<UUID, List<BlockingState>> entitlementStatesPerBundle = new HashMap<UUID, List<BlockingState>>();
+        // Optimization: build lookup tables for blocking states states
+        final Collection<BlockingState> accountBlockingStates = new LinkedList<BlockingState>();
+        final Map<UUID, List<BlockingState>> blockingStatesPerSubscription = new HashMap<UUID, List<BlockingState>>();
+        final Map<UUID, List<BlockingState>> blockingStatesPerBundle = new HashMap<UUID, List<BlockingState>>();
         for (final BlockingState blockingState : blockingStatesForAccount) {
-            if (!EntitlementService.ENTITLEMENT_SERVICE_NAME.equals(blockingState.getService())) {
-                continue;
-            } else if (BlockingStateType.SUBSCRIPTION.equals(blockingState.getType())) {
-                if (entitlementStatesPerSubscription.get(blockingState.getBlockedId()) == null) {
-                    entitlementStatesPerSubscription.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
+            if (BlockingStateType.SUBSCRIPTION.equals(blockingState.getType())) {
+                if (blockingStatesPerSubscription.get(blockingState.getBlockedId()) == null) {
+                    blockingStatesPerSubscription.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
                 }
-                entitlementStatesPerSubscription.get(blockingState.getBlockedId()).add(blockingState);
+                blockingStatesPerSubscription.get(blockingState.getBlockedId()).add(blockingState);
             } else if (BlockingStateType.SUBSCRIPTION_BUNDLE.equals(blockingState.getType())) {
-                if (entitlementStatesPerBundle.get(blockingState.getBlockedId()) == null) {
-                    entitlementStatesPerBundle.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
+                if (blockingStatesPerBundle.get(blockingState.getBlockedId()) == null) {
+                    blockingStatesPerBundle.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
                 }
-                entitlementStatesPerBundle.get(blockingState.getBlockedId()).add(blockingState);
+                blockingStatesPerBundle.get(blockingState.getBlockedId()).add(blockingState);
             } else if (BlockingStateType.ACCOUNT.equals(blockingState.getType()) &&
                        account.getId().equals(blockingState.getBlockedId())) {
-                accountEntitlementStates.add(blockingState);
+                accountBlockingStates.add(blockingState);
             }
         }
 
@@ -181,29 +179,39 @@ public class EventsStreamBuilder {
                                                                                                   return ProductCategory.BASE.equals(input.getLastActiveProduct().getCategory());
                                                                                               }
                                                                                           }).orNull();
-            final List<BlockingState> bundleEntitlementStates = Objects.firstNonNull(entitlementStatesPerBundle.get(bundleId), ImmutableList.<BlockingState>of());
+            final List<BlockingState> bundleBlockingStates = Objects.firstNonNull(blockingStatesPerBundle.get(bundleId), ImmutableList.<BlockingState>of());
 
             if (entitlementsPerBundle.get(bundleId) == null) {
                 entitlementsPerBundle.put(bundleId, new LinkedList<EventsStream>());
             }
 
             for (final SubscriptionBase subscription : allSubscriptionsForBundle) {
-                final List<BlockingState> subscriptionBlockingStatesOnDisk = Objects.firstNonNull(entitlementStatesPerSubscription.get(subscription.getId()), ImmutableList.<BlockingState>of());
+                final List<BlockingState> subscriptionBlockingStatesOnDisk = Objects.firstNonNull(blockingStatesPerSubscription.get(subscription.getId()), ImmutableList.<BlockingState>of());
 
-                // We cannot use blockingStatesForAccount here: we need subscriptionEntitlementStates to contain the events not on disk when building an EventsStream
+                // We cannot always use blockingStatesForAccount here: we need subscriptionBlockingStates to contain the events not on disk when building an EventsStream
                 // for an add-on - which means going through the magic of ProxyBlockingStateDao, which will recursively
                 // create EventsStream objects. To avoid an infinite recursion, bypass ProxyBlockingStateDao when it's not
                 // needed, i.e. if this EventStream is for a standalone or a base subscription
-                final List<BlockingState> subscriptionEntitlementStates = (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) ?
-                                                                          subscriptionBlockingStatesOnDisk :
-                                                                          blockingStateDao.getBlockingHistoryForService(subscriptionBlockingStatesOnDisk,
-                                                                                                                        bundle,
-                                                                                                                        baseSubscription,
-                                                                                                                        subscription,
-                                                                                                                        allSubscriptionsForBundle,
-                                                                                                                        internalTenantContext);
-
-                final EventsStream eventStream = buildForEntitlement(account, bundle, baseSubscription, subscription, allSubscriptionsForBundle, subscriptionEntitlementStates, bundleEntitlementStates, accountEntitlementStates, internalTenantContext);
+                final List<BlockingState> subscriptionBlockingStates;
+                if (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) {
+                    subscriptionBlockingStates = subscriptionBlockingStatesOnDisk;
+                } else {
+                    subscriptionBlockingStates = blockingStateDao.getBlockingHistory(subscriptionBlockingStatesOnDisk,
+                                                                                     bundle,
+                                                                                     baseSubscription,
+                                                                                     subscription,
+                                                                                     allSubscriptionsForBundle,
+                                                                                     internalTenantContext);
+
+                }
+
+                // Merge the BlockingStates
+                final Collection<BlockingState> blockingStateSet = new LinkedHashSet<BlockingState>(accountBlockingStates);
+                blockingStateSet.addAll(bundleBlockingStates);
+                blockingStateSet.addAll(subscriptionBlockingStates);
+                final List<BlockingState> blockingStates = BLOCKING_STATE_ORDERING.immutableSortedCopy(blockingStateSet);
+
+                final EventsStream eventStream = buildForEntitlement(account, bundle, baseSubscription, subscription, allSubscriptionsForBundle, blockingStates, internalTenantContext);
                 entitlementsPerBundle.get(bundleId).add(eventStream);
             }
         }
@@ -267,24 +275,56 @@ public class EventsStreamBuilder {
             throw new EntitlementApiException(e);
         }
 
-        final List<BlockingState> bundleEntitlementStates = BLOCKING_STATE_ORDERING.immutableSortedCopy(defaultBlockingStateDao.getBlockingHistoryForService(bundle.getId(), BlockingStateType.SUBSCRIPTION_BUNDLE, EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext));
-        final List<BlockingState> accountEntitlementStates = BLOCKING_STATE_ORDERING.immutableSortedCopy(defaultBlockingStateDao.getBlockingHistoryForService(account.getId(), BlockingStateType.ACCOUNT, EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext));
-        final ImmutableList<BlockingState> subscriptionEntitlementStatesOnDisk = BLOCKING_STATE_ORDERING.immutableSortedCopy(defaultBlockingStateDao.getBlockingHistoryForService(subscription.getId(), BlockingStateType.SUBSCRIPTION, EntitlementService.ENTITLEMENT_SERVICE_NAME, internalTenantContext));
+        // Retrieve the blocking states
+        final Collection<BlockingState> blockingStatesForAccount = defaultBlockingStateDao.getBlockingAllForAccountRecordId(internalTenantContext);
+
+        // Optimization: build lookup tables for blocking states states
+        final Collection<BlockingState> accountBlockingStates = new LinkedList<BlockingState>();
+        final Map<UUID, List<BlockingState>> blockingStatesPerSubscription = new HashMap<UUID, List<BlockingState>>();
+        final Map<UUID, List<BlockingState>> blockingStatesPerBundle = new HashMap<UUID, List<BlockingState>>();
+        for (final BlockingState blockingState : blockingStatesForAccount) {
+            if (BlockingStateType.SUBSCRIPTION.equals(blockingState.getType())) {
+                if (blockingStatesPerSubscription.get(blockingState.getBlockedId()) == null) {
+                    blockingStatesPerSubscription.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
+                }
+                blockingStatesPerSubscription.get(blockingState.getBlockedId()).add(blockingState);
+            } else if (BlockingStateType.SUBSCRIPTION_BUNDLE.equals(blockingState.getType())) {
+                if (blockingStatesPerBundle.get(blockingState.getBlockedId()) == null) {
+                    blockingStatesPerBundle.put(blockingState.getBlockedId(), new LinkedList<BlockingState>());
+                }
+                blockingStatesPerBundle.get(blockingState.getBlockedId()).add(blockingState);
+            } else if (BlockingStateType.ACCOUNT.equals(blockingState.getType()) &&
+                       account.getId().equals(blockingState.getBlockedId())) {
+                accountBlockingStates.add(blockingState);
+            }
+        }
+
+        final List<BlockingState> bundleBlockingStates = Objects.firstNonNull(blockingStatesPerBundle.get(subscription.getBundleId()), ImmutableList.<BlockingState>of());
+        final List<BlockingState> subscriptionBlockingStatesOnDisk = Objects.firstNonNull(blockingStatesPerSubscription.get(subscription.getId()), ImmutableList.<BlockingState>of());
 
-        // We need subscriptionEntitlementStates to contain the events not on disk when building an EventsStream
+        // We cannot always use blockingStatesForAccount here: we need subscriptionBlockingStates to contain the events not on disk when building an EventsStream
         // for an add-on - which means going through the magic of ProxyBlockingStateDao, which will recursively
         // create EventsStream objects. To avoid an infinite recursion, bypass ProxyBlockingStateDao when it's not
         // needed, i.e. if this EventStream is for a standalone or a base subscription
-        final List<BlockingState> subscriptionEntitlementStates = (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) ?
-                                                                  subscriptionEntitlementStatesOnDisk :
-                                                                  blockingStateDao.getBlockingHistoryForService(subscriptionEntitlementStatesOnDisk,
-                                                                                                                bundle,
-                                                                                                                baseSubscription,
-                                                                                                                subscription,
-                                                                                                                allSubscriptionsForBundle,
-                                                                                                                internalTenantContext);
-
-        return buildForEntitlement(account, bundle, baseSubscription, subscription, allSubscriptionsForBundle, subscriptionEntitlementStates, bundleEntitlementStates, accountEntitlementStates, internalTenantContext);
+        final Collection<BlockingState> subscriptionBlockingStates;
+        if (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) {
+            subscriptionBlockingStates = subscriptionBlockingStatesOnDisk;
+        } else {
+            subscriptionBlockingStates = blockingStateDao.getBlockingHistory(ImmutableList.<BlockingState>copyOf(subscriptionBlockingStatesOnDisk),
+                                                                             bundle,
+                                                                             baseSubscription,
+                                                                             subscription,
+                                                                             allSubscriptionsForBundle,
+                                                                             internalTenantContext);
+        }
+
+        // Merge the BlockingStates
+        final Collection<BlockingState> blockingStateSet = new LinkedHashSet<BlockingState>(accountBlockingStates);
+        blockingStateSet.addAll(bundleBlockingStates);
+        blockingStateSet.addAll(subscriptionBlockingStates);
+        final List<BlockingState> blockingStates = BLOCKING_STATE_ORDERING.immutableSortedCopy(blockingStateSet);
+
+        return buildForEntitlement(account, bundle, baseSubscription, subscription, allSubscriptionsForBundle, blockingStates, internalTenantContext);
     }
 
     private EventsStream buildForEntitlement(final Account account,
@@ -292,15 +332,11 @@ public class EventsStreamBuilder {
                                              @Nullable final SubscriptionBase baseSubscription,
                                              final SubscriptionBase subscription,
                                              final List<SubscriptionBase> allSubscriptionsForBundle,
-                                             final List<BlockingState> subscriptionEntitlementStates,
-                                             final List<BlockingState> bundleEntitlementStates,
-                                             final List<BlockingState> accountEntitlementStates,
+                                             final List<BlockingState> blockingStates,
                                              final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         return new DefaultEventsStream(account,
                                        bundle,
-                                       subscriptionEntitlementStates,
-                                       bundleEntitlementStates,
-                                       accountEntitlementStates,
+                                       blockingStates,
                                        checker,
                                        baseSubscription,
                                        subscription,
diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
index ea79a11..113fb50 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/api/TestDefaultSubscriptionApi.java
@@ -17,8 +17,10 @@
 package com.ning.billing.entitlement.api;
 
 import java.util.List;
+import java.util.UUID;
 
 import org.joda.time.LocalDate;
+import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import com.ning.billing.account.api.Account;
@@ -29,12 +31,50 @@ import com.ning.billing.catalog.api.PlanPhaseSpecifier;
 import com.ning.billing.catalog.api.PriceListSet;
 import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.entitlement.EntitlementTestSuiteWithEmbeddedDB;
+import com.ning.billing.junction.DefaultBlockingState;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNull;
 
 public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbeddedDB {
 
+    @Test(groups = "slow", description = "Verify blocking states are exposed in SubscriptionBundle")
+    public void testBlockingStatesInTimelineApi() throws Exception {
+        final LocalDate initialDate = new LocalDate(2013, 8, 7);
+        clock.setDay(initialDate);
+
+        final Account account = accountApi.createAccount(getAccountData(7), callContext);
+        final PlanPhaseSpecifier spec = new PlanPhaseSpecifier("Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
+        testListener.pushExpectedEvents(NextEvent.CREATE, NextEvent.CREATE, NextEvent.BLOCK);
+        final Entitlement entitlement1 = entitlementApi.createBaseEntitlement(account.getId(), spec, UUID.randomUUID().toString(), initialDate, callContext);
+        final Entitlement entitlement2 = entitlementApi.createBaseEntitlement(account.getId(), spec, UUID.randomUUID().toString(), initialDate, callContext);
+        entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(new DefaultBlockingState(account.getId(), BlockingStateType.ACCOUNT, "stateName", "service", false, false, false, clock.getUTCNow()),
+                                                                        internalCallContextFactory.createInternalCallContext(account.getId(), callContext));
+        assertListenerStatus();
+
+        final List<SubscriptionBundle> bundles = subscriptionApi.getSubscriptionBundlesForAccountId(account.getId(), callContext);
+        Assert.assertEquals(bundles.size(), 2);
+        // This will test the ordering as well
+        subscriptionBundleChecker(bundles, initialDate, entitlement1, 0);
+        subscriptionBundleChecker(bundles, initialDate, entitlement2, 1);
+    }
+
+    private void subscriptionBundleChecker(final List<SubscriptionBundle> bundles, final LocalDate initialDate, final Entitlement entitlement, final int idx) {
+        Assert.assertEquals(bundles.get(idx).getId(), entitlement.getBundleId());
+        Assert.assertEquals(bundles.get(idx).getSubscriptions().size(), 1);
+        Assert.assertEquals(bundles.get(idx).getSubscriptions().get(0).getId(), entitlement.getId());
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().size(), 4);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(0).getEffectiveDate(), initialDate);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(0).getSubscriptionEventType(), SubscriptionEventType.START_ENTITLEMENT);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(1).getEffectiveDate(), initialDate);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(1).getSubscriptionEventType(), SubscriptionEventType.START_BILLING);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(2).getEffectiveDate(), initialDate);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(2).getSubscriptionEventType(), SubscriptionEventType.SERVICE_STATE_CHANGE);
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(2).getServiceName(), "service");
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(2).getServiceStateName(), "stateName");
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(3).getEffectiveDate(), new LocalDate(2013, 9, 6));
+        Assert.assertEquals(bundles.get(idx).getTimeline().getSubscriptionEvents().get(3).getSubscriptionEventType(), SubscriptionEventType.PHASE);
+    }
 
     @Test(groups = "slow")
     public void testWithMultipleBundle() throws AccountApiException, SubscriptionApiException, EntitlementApiException {
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 78ba189..b623448 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
@@ -70,6 +70,9 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         clock.setDay(initialDate);
         account = accountApi.createAccount(getAccountData(7), callContext);
 
+        // Override the context with the right account record id
+        internalCallContext = internalCallContextFactory.createInternalCallContext(account.getId(), callContext);
+
         testListener.pushExpectedEvents(NextEvent.CREATE, NextEvent.CREATE);
 
         // Create base entitlement