killbill-memoizeit

Fix broken API getSubscriptionBundlesForAccountIdAndExternalKey

9/26/2013 6:59:30 PM

Details

diff --git a/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseInternalApi.java b/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseInternalApi.java
index 7696874..b8de4a1 100644
--- a/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseInternalApi.java
+++ b/api/src/main/java/com/ning/billing/subscription/api/SubscriptionBaseInternalApi.java
@@ -42,7 +42,7 @@ public interface SubscriptionBaseInternalApi {
     public SubscriptionBaseBundle createBundleForAccount(final UUID accountId, final String bundleName, final InternalCallContext context)
             throws SubscriptionBaseApiException;
 
-    public SubscriptionBaseBundle getBundleForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context)
+    public List<SubscriptionBaseBundle> getBundlesForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context)
             throws SubscriptionBaseApiException;
 
     public List<SubscriptionBaseBundle> getBundlesForAccount(final UUID accountId, final InternalTenantContext context);
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 c895e02..1452b63 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
@@ -16,6 +16,7 @@
 
 package com.ning.billing.entitlement.api;
 
+import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -180,8 +181,9 @@ public class DefaultEntitlementApi implements EntitlementApi {
 
             final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
 
-            final LocalDate entitlementEffectiveEndDate = getEffectiveEndDate(bundle.getAccountId(), subscription, account.getTimeZone(), context);
-            final EntitlementState entitlementState = getStateForEntitlement(entitlementEffectiveEndDate, subscription, account.getTimeZone(), context);
+            final InternalTenantContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalTenantContext(account.getId(), tenantContext);
+            final LocalDate entitlementEffectiveEndDate = getEffectiveEndDate(bundle.getAccountId(), subscription, account.getTimeZone(), contextWithValidAccountRecordId);
+            final EntitlementState entitlementState = getStateForEntitlement(entitlementEffectiveEndDate, subscription, account.getTimeZone(), contextWithValidAccountRecordId);
 
 
             return new DefaultEntitlement(dateHelper, subscription, bundle.getAccountId(), bundle.getExternalKey(), entitlementState, entitlementEffectiveEndDate, account.getTimeZone(),
@@ -199,7 +201,9 @@ public class DefaultEntitlementApi implements EntitlementApi {
         try {
             final SubscriptionBaseBundle bundle = subscriptionInternalApi.getBundleFromId(bundleId, context);
             final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
-            return getAllEntitlementsForBundleId(bundleId, bundle.getAccountId(), account.getTimeZone(), bundle.getExternalKey(), context);
+
+            final InternalTenantContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalTenantContext(account.getId(), tenantContext);
+            return getAllEntitlementsForBundleId(bundleId, bundle.getAccountId(), account.getTimeZone(), bundle.getExternalKey(), contextWithValidAccountRecordId);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         } catch (AccountApiException e) {
@@ -209,11 +213,15 @@ public class DefaultEntitlementApi implements EntitlementApi {
 
     @Override
     public List<Entitlement> getAllEntitlementsForAccountIdAndExternalKey(final UUID accountId, final String externalKey, final TenantContext tenantContext) throws EntitlementApiException {
-        final InternalTenantContext context = internalCallContextFactory.createInternalTenantContext(tenantContext);
+        final InternalTenantContext context = internalCallContextFactory.createInternalTenantContext(accountId, tenantContext);
         try {
-            final SubscriptionBaseBundle bundle = subscriptionInternalApi.getBundleForAccountAndKey(accountId, externalKey, context);
-            final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
-            return getAllEntitlementsForBundleId(bundle.getId(), bundle.getAccountId(), account.getTimeZone(), bundle.getExternalKey(), context);
+            final List<SubscriptionBaseBundle> bundles = subscriptionInternalApi.getBundlesForAccountAndKey(accountId, externalKey, context);
+            if (bundles.size() == 0) {
+                return Collections.emptyList();
+            }
+            final Account account = accountApi.getAccountById(bundles.get(0).getAccountId(), context);
+            return getEntitlementsForBundles(bundles, account, context);
+
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         } catch (AccountApiException e) {
@@ -230,8 +238,11 @@ public class DefaultEntitlementApi implements EntitlementApi {
         } catch (AccountApiException e) {
             throw new EntitlementApiException(e);
         }
+        final List<SubscriptionBaseBundle> bundles = subscriptionInternalApi.getBundlesForAccount(account.getId(), context);
+        return getEntitlementsForBundles(bundles, account, context);
+    }
 
-        final List<SubscriptionBaseBundle> bundles = subscriptionInternalApi.getBundlesForAccount(accountId, context);
+    private List<Entitlement> getEntitlementsForBundles(final List<SubscriptionBaseBundle> bundles, final Account account, final InternalTenantContext context) {
         final Map<UUID, List<SubscriptionBase>> subscriptionsPerBundle = subscriptionInternalApi.getSubscriptionsForAccount(context);
         final List<Entitlement> result = new LinkedList<Entitlement>();
         for (final SubscriptionBaseBundle bundle : bundles) {
@@ -425,7 +436,10 @@ public class DefaultEntitlementApi implements EntitlementApi {
 
         final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(sourceAccountId, context);
         try {
-            final SubscriptionBaseBundle bundle = subscriptionInternalApi.getBundleForAccountAndKey(sourceAccountId, externalKey, contextWithValidAccountRecordId);
+            final SubscriptionBaseBundle bundle = subscriptionInternalApi.getActiveBundleForKey(externalKey, contextWithValidAccountRecordId);
+            if (!bundle.getAccountId().equals(sourceAccountId)) {
+                throw new EntitlementApiException(new SubscriptionBaseApiException(ErrorCode.SUB_GET_INVALID_BUNDLE_KEY, externalKey));
+            }
             final SubscriptionBase baseSubscription = subscriptionInternalApi.getBaseSubscription(bundle.getId(), contextWithValidAccountRecordId);
 
             final DateTime requestedDate = dateHelper.fromLocalDateAndReferenceTime(effectiveDate, baseSubscription.getStartDate(), contextWithValidAccountRecordId);
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 9bd705f..d6b01c3 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
@@ -104,24 +104,15 @@ public class DefaultSubscriptionApi implements SubscriptionApi {
 
 
     @Override
-    public SubscriptionBundle getSubscriptionBundleForAccountIdAndExternalKey(final UUID accountId, final String externalKey, final TenantContext context) throws SubscriptionApiException {
+    public List<SubscriptionBundle> getSubscriptionBundlesForAccountIdAndExternalKey(final UUID accountId, final String externalKey, final TenantContext context) throws SubscriptionApiException {
 
         try {
-            // If active bundle exists on this account, this is what we are looking for.
-            try {
-                final SubscriptionBundle activeBundleForkey = getActiveSubscriptionBundleForExternalKey(externalKey, context);
-                if (activeBundleForkey.getAccountId().equals(accountId)) {
-                    return activeBundleForkey;
-                }
-            } catch (SubscriptionApiException ignore) {
-            }
 
-            // If not return first cancelled bundle if this is exists
             final List<Entitlement> entitlements = entitlementApi.getAllEntitlementsForAccountIdAndExternalKey(accountId, externalKey, context);
             if (entitlements.isEmpty()) {
                 throw new SubscriptionApiException(ErrorCode.SUB_GET_INVALID_BUNDLE_KEY, externalKey);
             }
-            return getSubscriptionBundleFromEntitlements(entitlements.get(0).getBundleId(), entitlements, context);
+            return getSubscriptionBundles(entitlements, context);
         } catch (EntitlementApiException e) {
             throw new SubscriptionApiException(e);
         } catch (SubscriptionBaseApiException e) {
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 8013937..3ac0011 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
@@ -93,6 +93,8 @@ public class TestDefaultSubscriptionApi extends EntitlementTestSuiteWithEmbedded
             assertNull(secondbundle.getSubscriptions().get(0).getBillingEndDate());
             assertEquals(secondbundle.getOriginalCreatedDate().compareTo(firstbundle.getCreatedDate()), 0);
 
+            final List<SubscriptionBundle> bundles2Again = subscriptionApi.getSubscriptionBundlesForAccountIdAndExternalKey(account.getId(), externalKey, callContext);
+            assertEquals(bundles2Again.size(), 2);
 
             clock.addDays(3);
             final Account account2 = accountApi.createAccount(getAccountData(7), callContext);
diff --git a/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java b/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
index 8538783..eb26739 100644
--- a/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
+++ b/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
@@ -114,6 +114,7 @@ public class AccountResource extends JaxRsResourceBase {
     private final InvoicePaymentApi invoicePaymentApi;
     private final PaymentApi paymentApi;
     private final OverdueUserApi overdueApi;
+
     @Inject
     public AccountResource(final JaxrsUriBuilder uriBuilder,
                            final AccountUserApi accountApi,
@@ -177,21 +178,18 @@ public class AccountResource extends JaxRsResourceBase {
 
         final UUID uuid = UUID.fromString(accountId);
         accountUserApi.getAccountById(uuid, tenantContext);
-        if (externalKey != null) {
 
-            final SubscriptionBundle bundle = subscriptionApi.getSubscriptionBundleForAccountIdAndExternalKey(uuid, externalKey, tenantContext);
-            final BundleJson json = new BundleJson(bundle, null, null, null);
-            return Response.status(Status.OK).entity(json).build();
-        } else {
-            final List<SubscriptionBundle> bundles = subscriptionApi.getSubscriptionBundlesForAccountId(uuid, tenantContext);
-            final Collection<BundleJson> result = Collections2.transform(bundles, new Function<SubscriptionBundle, BundleJson>() {
-                @Override
-                public BundleJson apply(final SubscriptionBundle input) {
-                    return new BundleJson(input, null, null, null);
-                }
-            });
-            return Response.status(Status.OK).entity(result).build();
-        }
+        final List<SubscriptionBundle> bundles = (externalKey != null) ?
+                                                 subscriptionApi.getSubscriptionBundlesForAccountIdAndExternalKey(uuid, externalKey, tenantContext) :
+                                                 subscriptionApi.getSubscriptionBundlesForAccountId(uuid, tenantContext);
+
+        final Collection<BundleJson> result = Collections2.transform(bundles, new Function<SubscriptionBundle, BundleJson>() {
+            @Override
+            public BundleJson apply(final SubscriptionBundle input) {
+                return new BundleJson(input, null, null, null);
+            }
+        });
+        return Response.status(Status.OK).entity(result).build();
     }
 
     @GET
@@ -494,7 +492,7 @@ public class AccountResource extends JaxRsResourceBase {
     @Path("/{accountId:" + UUID_PATTERN + "}/" + CHARGEBACKS)
     @Produces(APPLICATION_JSON)
     public Response getChargebacksForAccount(@PathParam("accountId") final String accountIdStr,
-                                  @javax.ws.rs.core.Context final HttpServletRequest request) {
+                                             @javax.ws.rs.core.Context final HttpServletRequest request) {
         final UUID accountId = UUID.fromString(accountIdStr);
         final List<InvoicePayment> chargebacks = invoicePaymentApi.getChargebacksByAccountId(accountId, context.createContext(request));
         final List<ChargebackJson> chargebacksJson = new ArrayList<ChargebackJson>();
@@ -529,7 +527,6 @@ public class AccountResource extends JaxRsResourceBase {
     }
 
 
-
     /*
      * ************************** OVERDUE ********************************
      */
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java b/subscription/src/main/java/com/ning/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
index da6a179..cebd98a 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
@@ -168,13 +168,9 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
     }
 
     @Override
-    public SubscriptionBaseBundle getBundleForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) throws SubscriptionBaseApiException {
-        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = dao.getSubscriptionBundlesFromAccountAndKey(accountId, bundleKey, context);
-        final SubscriptionBaseBundle result = getActiveBundleForKeyNotException(bundlesForAccountAndKey, dao, clock, context);
-        if (result == null) {
-            throw new SubscriptionBaseApiException(ErrorCode.SUB_GET_INVALID_BUNDLE_KEY, bundleKey);
-        }
-        return result;
+    public List<SubscriptionBaseBundle> getBundlesForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) throws SubscriptionBaseApiException {
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = dao.getSubscriptionBundlesForAccountAndKey(accountId, bundleKey, context);
+        return bundlesForAccountAndKey;
     }
 
     @Override
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/timeline/DefaultSubscriptionBaseTimelineApi.java b/subscription/src/main/java/com/ning/billing/subscription/api/timeline/DefaultSubscriptionBaseTimelineApi.java
index 7eae1db..06ae6f4 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/timeline/DefaultSubscriptionBaseTimelineApi.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/timeline/DefaultSubscriptionBaseTimelineApi.java
@@ -98,7 +98,7 @@ public class DefaultSubscriptionBaseTimelineApi extends SubscriptionApiBase impl
     @Override
     public BundleBaseTimeline getBundleTimeline(final UUID accountId, final String bundleName, final TenantContext context)
             throws SubscriptionBaseRepairException {
-        final List<SubscriptionBaseBundle> bundles = dao.getSubscriptionBundlesFromAccountAndKey(accountId, bundleName, internalCallContextFactory.createInternalTenantContext(context));
+        final List<SubscriptionBaseBundle> bundles = dao.getSubscriptionBundlesForAccountAndKey(accountId, bundleName, internalCallContextFactory.createInternalTenantContext(context));
         final SubscriptionBaseBundle bundle = bundles.size() > 0 ? bundles.get(bundles.size() - 1) : null;
         return getBundleTimelineInternal(bundle, bundleName + " [accountId= " + accountId.toString() + "]", context);
     }
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/transfer/DefaultSubscriptionBaseTransferApi.java b/subscription/src/main/java/com/ning/billing/subscription/api/transfer/DefaultSubscriptionBaseTransferApi.java
index 60b4926..7e62ee8 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/transfer/DefaultSubscriptionBaseTransferApi.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/transfer/DefaultSubscriptionBaseTransferApi.java
@@ -201,7 +201,7 @@ public class DefaultSubscriptionBaseTransferApi extends SubscriptionApiBase impl
                 throw new SubscriptionBaseTransferApiException(ErrorCode.SUB_TRANSFER_INVALID_EFF_DATE, effectiveTransferDate);
             }
 
-            final List<SubscriptionBaseBundle> bundlesForAccountAndKey = dao.getSubscriptionBundlesFromAccountAndKey(sourceAccountId, bundleKey, fromInternalCallContext);
+            final List<SubscriptionBaseBundle> bundlesForAccountAndKey = dao.getSubscriptionBundlesForAccountAndKey(sourceAccountId, bundleKey, fromInternalCallContext);
             final SubscriptionBaseBundle bundle = DefaultSubscriptionInternalApi.getActiveBundleForKeyNotException(bundlesForAccountAndKey, dao, clock, fromInternalCallContext);
             if (bundle == null) {
                 throw new SubscriptionBaseTransferApiException(ErrorCode.SUB_CREATE_NO_BUNDLE, bundleKey);
diff --git a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/DefaultSubscriptionDao.java b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/DefaultSubscriptionDao.java
index 549b5b5..7aec7aa 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/DefaultSubscriptionDao.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/DefaultSubscriptionDao.java
@@ -118,7 +118,7 @@ public class DefaultSubscriptionDao implements SubscriptionDao {
     }
 
     @Override
-    public List<SubscriptionBaseBundle> getSubscriptionBundlesFromAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
+    public List<SubscriptionBaseBundle> getSubscriptionBundlesForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
         return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<List<SubscriptionBaseBundle>>() {
             @Override
             public List<SubscriptionBaseBundle> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
diff --git a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/RepairSubscriptionDao.java b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/RepairSubscriptionDao.java
index 3f6146c..b05e650 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/RepairSubscriptionDao.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/RepairSubscriptionDao.java
@@ -306,7 +306,7 @@ public class RepairSubscriptionDao implements SubscriptionDao, RepairSubscriptio
     }
 
     @Override
-    public List<SubscriptionBaseBundle> getSubscriptionBundlesFromAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
+    public List<SubscriptionBaseBundle> getSubscriptionBundlesForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
         throw new SubscriptionBaseError(NOT_IMPLEMENTED);
     }
 }
diff --git a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/SubscriptionDao.java b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/SubscriptionDao.java
index 722c031..b7cffc0 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/engine/dao/SubscriptionDao.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/engine/dao/SubscriptionDao.java
@@ -39,7 +39,7 @@ public interface SubscriptionDao {
 
     public List<SubscriptionBaseBundle> getSubscriptionBundlesForKey(String bundleKey, InternalTenantContext context);
 
-    public List<SubscriptionBaseBundle> getSubscriptionBundlesFromAccountAndKey(UUID accountId, String bundleKey, InternalTenantContext context);
+    public List<SubscriptionBaseBundle> getSubscriptionBundlesForAccountAndKey(UUID accountId, String bundleKey, InternalTenantContext context);
 
     public SubscriptionBaseBundle getSubscriptionBundleFromId(UUID bundleId, InternalTenantContext context);
 
diff --git a/subscription/src/test/java/com/ning/billing/subscription/api/transfer/TestTransfer.java b/subscription/src/test/java/com/ning/billing/subscription/api/transfer/TestTransfer.java
index 44f8256..e6dd7ff 100644
--- a/subscription/src/test/java/com/ning/billing/subscription/api/transfer/TestTransfer.java
+++ b/subscription/src/test/java/com/ning/billing/subscription/api/transfer/TestTransfer.java
@@ -149,8 +149,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         assertTrue(oldBaseSubscription.getEndDate().compareTo(transferRequestedDate) == 0);
 
         // CHECK NEW BUNDLE EXIST, WITH ONE SUBSCRIPTION STARTING ON TRANSFER_DATE
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
 
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
 
@@ -198,7 +200,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         assertTrue(oldBaseSubscription.getFutureEndDate().compareTo(ctd) == 0);
 
         // CHECK NEW BUNDLE EXIST, WITH ONE SUBSCRIPTION STARTING ON TRANSFER_DATE
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
+
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
 
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
@@ -248,8 +253,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         assertTrue(oldBaseSubscription.getEndDate().compareTo(transferRequestedDate) == 0);
 
         // CHECK NEW BUNDLE EXIST, WITH ONE SUBSCRIPTION STARTING ON TRANSFER_DATE
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
 
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
 
@@ -297,8 +304,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         assertTrue(oldBaseSubscription.getFutureEndDate().compareTo(ctd) == 0);
 
         // CHECK NEW BUNDLE EXIST, WITH ONE SUBSCRIPTION STARTING ON TRANSFER_DATE
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
 
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
 
@@ -389,7 +398,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         assertTrue(testListener.isCompleted(3000));
 
         // RETRIEVE NEW BUNDLE AND CHECK SUBSCRIPTIONS
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
+
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 3);
         boolean foundBP = false;
@@ -479,7 +491,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         transferApi.transferBundle(bundle.getAccountId(), newAccountId, bundle.getExternalKey(), transferRequestedDate, true, false, callContext);
         assertTrue(testListener.isCompleted(3000));
 
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
+
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
     }
@@ -522,7 +537,10 @@ public class TestTransfer extends SubscriptionTestSuiteWithEmbeddedDB {
         transferApi.transferBundle(bundle.getAccountId(), newAccountId, bundle.getExternalKey(), transferRequestedDate, true, false, callContext);
         assertTrue(testListener.isCompleted(3000));
 
-        final SubscriptionBaseBundle newBundle = subscriptionInternalApi.getBundleForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        final List<SubscriptionBaseBundle> bundlesForAccountAndKey = subscriptionInternalApi.getBundlesForAccountAndKey(newAccountId, bundle.getExternalKey(), internalCallContext);
+        assertEquals(bundlesForAccountAndKey.size(), 1);
+
+        final SubscriptionBaseBundle newBundle = bundlesForAccountAndKey.get(0);
         final List<SubscriptionBase> subscriptions = subscriptionInternalApi.getSubscriptionsForBundle(newBundle.getId(), internalCallContext);
         assertEquals(subscriptions.size(), 1);
     }
diff --git a/subscription/src/test/java/com/ning/billing/subscription/engine/dao/MockSubscriptionDaoMemory.java b/subscription/src/test/java/com/ning/billing/subscription/engine/dao/MockSubscriptionDaoMemory.java
index 99b32b9..f714b32 100644
--- a/subscription/src/test/java/com/ning/billing/subscription/engine/dao/MockSubscriptionDaoMemory.java
+++ b/subscription/src/test/java/com/ning/billing/subscription/engine/dao/MockSubscriptionDaoMemory.java
@@ -124,7 +124,7 @@ public class MockSubscriptionDaoMemory implements SubscriptionDao {
     }
 
     @Override
-    public List<SubscriptionBaseBundle> getSubscriptionBundlesFromAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
+    public List<SubscriptionBaseBundle> getSubscriptionBundlesForAccountAndKey(final UUID accountId, final String bundleKey, final InternalTenantContext context) {
         final List<SubscriptionBaseBundle> results = new ArrayList<SubscriptionBaseBundle>();
         for (final SubscriptionBaseBundle cur : bundles) {
             if (cur.getExternalKey().equals(bundleKey) && cur.getAccountId().equals(accountId)) {