killbill-memoizeit

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
index e0f3b60..9150555 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
@@ -226,10 +226,15 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
 
             final Collection<BlockingState> blockingStatesNotOnDisk = eventsStream.computeAddonsBlockingStatesForFutureSubscriptionBaseEvents();
 
-            // Inject the extra blocking states into the stream
+            // Inject the extra blocking states into the stream if needed
             for (final BlockingState blockingState : blockingStatesNotOnDisk) {
-                final BlockingStateModelDao blockingStateModelDao = new BlockingStateModelDao(blockingState, now, now);
-                blockingStatesOnDiskCopy.add(BlockingStateModelDao.toBlockingState(blockingStateModelDao));
+                // In case we're coming from getBlockingHistoryForService / getBlockingAll, make sure we don't add
+                // blocking states for other add-ons on that base subscription
+                if (blockingStateType == null ||
+                    (BlockingStateType.SUBSCRIPTION.equals(blockingStateType) && blockingState.getBlockedId().equals(blockableId))) {
+                    final BlockingStateModelDao blockingStateModelDao = new BlockingStateModelDao(blockingState, now, now);
+                    blockingStatesOnDiskCopy.add(BlockingStateModelDao.toBlockingState(blockingStateModelDao));
+                }
             }
         }
 

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 bb5f2d9..5398bce 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
@@ -69,7 +69,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         baseEntitlement = (DefaultEntitlement) entitlementApi.createBaseEntitlement(account.getId(), baseSpec, account.getExternalKey(), initialDate, callContext);
 
         // Add ADD_ON
-        final PlanPhaseSpecifier addOnSpec = new PlanPhaseSpecifier("Telescopic-Scope", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
+        final PlanPhaseSpecifier addOnSpec = new PlanPhaseSpecifier("Telescopic-Scope", ProductCategory.ADD_ON, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         addOnEntitlement = (DefaultEntitlement) entitlementApi.addEntitlement(baseEntitlement.getBundleId(), addOnSpec, initialDate, callContext);
 
         // Verify the initial state
@@ -227,6 +227,26 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         checkBlockingStatesDAO(changedBaseEntitlement, cancelledAddOnEntitlement, baseEffectiveCancellationOrChangeDate, false);
     }
 
+    @Test(groups = "slow", description = "Verify we don't mix add-ons for EOT changes")
+    public void testChangePlanEOTWith2AddOns() throws Exception {
+        // Add a second ADD_ON (Laser-Scope is available, not included)
+        testListener.pushExpectedEvents(NextEvent.CREATE);
+        final PlanPhaseSpecifier secondAddOnSpec = new PlanPhaseSpecifier("Laser-Scope", ProductCategory.ADD_ON, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
+        final DefaultEntitlement secondAddOnEntitlement = (DefaultEntitlement) entitlementApi.addEntitlement(baseEntitlement.getBundleId(), secondAddOnSpec, clock.getUTCToday(), callContext);
+        assertListenerStatus();
+
+        // Change plan EOT to Assault-Rifle (Telescopic-Scope is included)
+        final DefaultEntitlement changedBaseEntitlement = (DefaultEntitlement) baseEntitlement.changePlanWithDate("Assault-Rifle", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, new LocalDate(2013, 10, 7), callContext);
+        // No blocking event (EOT)
+        assertListenerStatus();
+
+        // Verify the blocking states DAO adds events not on disk for the first add-on...
+        checkBlockingStatesDAO(changedBaseEntitlement, addOnEntitlement, baseEffectiveCancellationOrChangeDate, false);
+        // ...but not for the second one
+        final List<BlockingState> blockingStatesForSecondAddOn = blockingStateDao.getBlockingAll(secondAddOnEntitlement.getId(), BlockingStateType.SUBSCRIPTION, internalCallContext);
+        Assert.assertEquals(blockingStatesForSecondAddOn.size(), 0);
+    }
+
     @Test(groups = "slow", description = "Verify add-ons blocking states are added for IMM change plans")
     public void testChangePlanIMM() throws Exception {
         // Approximate check, as the blocking state check (checkBlockingStatesDAO) could be a bit off