killbill-memoizeit

Details

diff --git a/api/src/main/java/com/ning/billing/subscription/api/user/EntitlementStatusDryRun.java b/api/src/main/java/com/ning/billing/subscription/api/user/EntitlementStatusDryRun.java
new file mode 100644
index 0000000..5bca65e
--- /dev/null
+++ b/api/src/main/java/com/ning/billing/subscription/api/user/EntitlementStatusDryRun.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2010-2013 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.subscription.api.user;
+
+import java.util.UUID;
+
+import com.ning.billing.catalog.api.BillingPeriod;
+import com.ning.billing.catalog.api.PhaseType;
+
+public interface EntitlementStatusDryRun {
+
+    public UUID getId();
+
+    public String getProductName();
+
+    public BillingPeriod getBillingPeriod();
+
+    public String getPriceList();
+
+    public PhaseType getPhaseType();
+
+    public DryRunChangeReason getReason();
+
+    public enum DryRunChangeReason {
+        AO_INCLUDED_IN_NEW_PLAN,
+        AO_NOT_AVAILABLE_IN_NEW_PLAN,
+        AO_AVAILABLE_IN_NEW_PLAN
+    }
+}
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 407f104..ef18662 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
@@ -45,7 +45,6 @@ import com.ning.billing.entitlement.dao.BlockingStateDao;
 import com.ning.billing.subscription.api.SubscriptionBase;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
-import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
@@ -54,7 +53,6 @@ import com.ning.billing.util.callcontext.TenantContext;
 import com.ning.billing.util.svcapi.account.AccountInternalApi;
 import com.ning.billing.util.svcapi.junction.DefaultBlockingState;
 import com.ning.billing.util.svcapi.subscription.SubscriptionBaseInternalApi;
-import com.ning.billing.util.timezone.DateAndTimeZoneContext;
 
 import com.google.common.base.Function;
 import com.google.common.collect.Collections2;
@@ -115,8 +113,10 @@ public class DefaultEntitlementApi implements EntitlementApi {
             final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(bundle.getAccountId(), callContext);
             final BlockingState currentBaseState =  blockingStateDao.getBlockingStateForService(baseSubscription.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, contextWithValidAccountRecordId);
 
+            final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
+
             // Check if there is a BP and if it is active
-            final EntitlementState baseEntitlementState = getStateForEntitlement(baseSubscription, currentBaseState, contextWithValidAccountRecordId);
+            final EntitlementState baseEntitlementState = getStateForEntitlement(baseSubscription, currentBaseState, account.getTimeZone(), contextWithValidAccountRecordId);
             if (baseSubscription.getCategory() != ProductCategory.BASE ||
                 baseEntitlementState != EntitlementState.ACTIVE) {
                 throw new EntitlementApiException(ErrorCode.SUB_GET_NO_SUCH_BASE_SUBSCRIPTION, baseSubscription.getBundleId());
@@ -128,11 +128,10 @@ public class DefaultEntitlementApi implements EntitlementApi {
                 throw new EntitlementApiException(new BlockingApiException(ErrorCode.BLOCK_BLOCKED_ACTION, BlockingChecker.ACTION_CHANGE, BlockingChecker.TYPE_SUBSCRIPTION, baseSubscription.getId().toString()));
             }
 
-            final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
 
             final DateTime requestedDate = dateHelper.fromNowAndReferenceTime(baseSubscription.getStartDate(), contextWithValidAccountRecordId);
             final SubscriptionBase subscription = subscriptionInternalApi.createSubscription(baseSubscription.getBundleId(), planPhaseSpecifier, requestedDate, context);
-            return new DefaultEntitlement(dateHelper, subscription, bundle.getAccountId(), bundle.getExternalKey(), getStateForEntitlement(subscription, currentBaseState, context), null, account.getTimeZone(),
+            return new DefaultEntitlement(dateHelper, subscription, bundle.getAccountId(), bundle.getExternalKey(), getStateForEntitlement(subscription, currentBaseState, account.getTimeZone(), context), null, account.getTimeZone(),
                                           internalCallContextFactory, blockingStateDao, clock, checker);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
@@ -153,7 +152,7 @@ public class DefaultEntitlementApi implements EntitlementApi {
             final Account account = accountApi.getAccountById(bundle.getAccountId(), context);
             final BlockingState currentState =  blockingStateDao.getBlockingStateForService(subscription.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, context);
 
-            return new DefaultEntitlement(dateHelper, subscription, bundle.getAccountId(), bundle.getExternalKey(), getStateForEntitlement(subscription, currentState, context), currentState, account.getTimeZone(),
+            return new DefaultEntitlement(dateHelper, subscription, bundle.getAccountId(), bundle.getExternalKey(), getStateForEntitlement(subscription, currentState, account.getTimeZone(), context), currentState, account.getTimeZone(),
                                           internalCallContextFactory, blockingStateDao, clock, checker);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
@@ -210,7 +209,7 @@ public class DefaultEntitlementApi implements EntitlementApi {
                     final BlockingState currentState =  blockingStateDao.getBlockingStateForService(input.getId(), EntitlementService.ENTITLEMENT_SERVICE_NAME, context);
 
                     return new DefaultEntitlement(dateHelper, input, accountId, externalKey,
-                                                  getStateForEntitlement(input, currentState, context),
+                                                  getStateForEntitlement(input, currentState, account.getTimeZone(), context),
                                                   currentState,
                                                   account.getTimeZone(),
                                                   internalCallContextFactory, blockingStateDao, clock, checker);
@@ -221,10 +220,13 @@ public class DefaultEntitlementApi implements EntitlementApi {
         }
     }
 
-    private EntitlementState getStateForEntitlement(final SubscriptionBase subscriptionBase, final BlockingState currentState, final InternalTenantContext context) {
+
+    private EntitlementState getStateForEntitlement(final SubscriptionBase subscriptionBase, final BlockingState currentState, final DateTimeZone accountTimeZone, final InternalTenantContext context) {
 
         // Current state for the ENTITLEMENT_SERVICE_NAME is set to cancelled
-        if (currentState != null && currentState.getStateName().equals(ENT_STATE_CANCELLED)) {
+        if (currentState != null &&
+            currentState.getStateName().equals(ENT_STATE_CANCELLED) &&
+                dateHelper.isBeforeOrEqualsNow(currentState.getEffectiveDate(), accountTimeZone)) {
             return EntitlementState.CANCELLED;
         }
 
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/EntitlementDateHelper.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/EntitlementDateHelper.java
index e1a1697..573c870 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/EntitlementDateHelper.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/EntitlementDateHelper.java
@@ -9,7 +9,6 @@ import com.ning.billing.account.api.AccountApiException;
 import com.ning.billing.clock.Clock;
 import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.svcapi.account.AccountInternalApi;
-import com.ning.billing.util.timezone.DateAndTimeZoneContext;
 
 public class EntitlementDateHelper {
 
@@ -21,52 +20,53 @@ public class EntitlementDateHelper {
         this.clock = clock;
     }
 
-    /**
-     * Returns a DateTime that is equals or beforeNow and whose LocalDate using the account timeZone is the one provided
-     * <p/>
-     * Relies on the subscriptionStartDate for the reference time
-     *
-     * @param requestedDate
-     * @param referenceDateTime
-     * @param callContext
-     * @return
-     * @throws EntitlementApiException
-     */
+
     public DateTime fromLocalDateAndReferenceTime(final LocalDate requestedDate, final DateTime referenceDateTime, final InternalCallContext callContext) throws EntitlementApiException {
         try {
-            /*
-            final Account account = accountApi.getAccountByRecordId(callContext.getAccountRecordId(), callContext);
-            final DateAndTimeZoneContext timeZoneContext = new DateAndTimeZoneContext(subscriptionStartDate, account.getTimeZone(), clock);
-            final DateTime computedTime = timeZoneContext.computeUTCDateTimeFromLocalDate(requestedDate);
-
-            return computedTime;
-            */
 
             final Account account = accountApi.getAccountByRecordId(callContext.getAccountRecordId(), callContext);
-            final LocalDate localDateNowInAccountTimezone = new LocalDate(requestedDate, account.getTimeZone());
-
-            // Datetime from local date in account timezone and with given reference time
-            final DateTime t1 = localDateNowInAccountTimezone.toDateTime(referenceDateTime.toLocalTime(), account.getTimeZone());
-            // Datetime converted back in UTC
-            final DateTime t2 = new DateTime(t1, DateTimeZone.UTC);
-            return t2;
-
+            return fromLocalDateAndReferenceTime(requestedDate, referenceDateTime, account.getTimeZone());
         } catch (AccountApiException e) {
             throw new EntitlementApiException(e);
         }
     }
 
+    public DateTime fromLocalDateAndReferenceTime(final LocalDate requestedDate, final DateTime referenceDateTime, final DateTimeZone accountTimeZone) throws EntitlementApiException {
+        final LocalDate localDateNowInAccountTimezone = new LocalDate(requestedDate, accountTimeZone);
+
+        // Datetime from local date in account timezone and with given reference time
+        final DateTime t1 = localDateNowInAccountTimezone.toDateTime(referenceDateTime.toLocalTime(), accountTimeZone);
+        // Datetime converted back in UTC
+        final DateTime t2 = new DateTime(t1, DateTimeZone.UTC);
+        return t2;
+
+    }
+
+
     public DateTime fromNowAndReferenceTime(final DateTime referenceDateTime, final InternalCallContext callContext) throws EntitlementApiException {
         try {
             final Account account = accountApi.getAccountByRecordId(callContext.getAccountRecordId(), callContext);
-            final LocalDate localDateNowInAccountTimezone = new LocalDate(clock.getUTCNow(), account.getTimeZone());
-            // Datetime from local date in account timezone and with given reference time
-            final DateTime t1 = localDateNowInAccountTimezone.toDateTime(referenceDateTime.toLocalTime(), account.getTimeZone());
-            // Datetime converted back in UTC
-            final DateTime t2 = new DateTime(t1, DateTimeZone.UTC);
-            return t2;
+            return fromNowAndReferenceTime(referenceDateTime, account.getTimeZone());
         } catch (AccountApiException e) {
             throw new EntitlementApiException(e);
         }
     }
+
+    public DateTime fromNowAndReferenceTime(final DateTime referenceDateTime, final DateTimeZone accountTimeZone) {
+        final LocalDate localDateNowInAccountTimezone = new LocalDate(clock.getUTCNow(), accountTimeZone);
+        // Datetime from local date in account timezone and with given reference time
+        final DateTime t1 = localDateNowInAccountTimezone.toDateTime(referenceDateTime.toLocalTime(), accountTimeZone);
+        // Datetime converted back in UTC
+        final DateTime t2 = new DateTime(t1, DateTimeZone.UTC);
+        return t2;
+    }
+
+    // STEPH_ENT test
+    public boolean isBeforeOrEqualsNow(final DateTime inputDate, final DateTimeZone accountTimeZone) {
+
+        final LocalDate localDateNowInAccountTimezone = new LocalDate(clock.getUTCNow(), accountTimeZone);
+        final LocalDate targetDateInAccountTimezone = new LocalDate(inputDate, accountTimeZone);
+
+        return targetDateInAccountTimezone.compareTo(localDateNowInAccountTimezone) <= 0;
+    }
 }
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 03a888f..60bc2fd 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
@@ -39,20 +39,19 @@ import com.ning.billing.clock.Clock;
 import com.ning.billing.clock.DefaultClock;
 import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 import com.ning.billing.subscription.api.SubscriptionApiBase;
+import com.ning.billing.subscription.api.SubscriptionBase;
 import com.ning.billing.subscription.api.user.DefaultEffectiveSubscriptionEvent;
 import com.ning.billing.subscription.api.user.DefaultSubscriptionBase;
 import com.ning.billing.subscription.api.user.DefaultSubscriptionBaseApiService;
 import com.ning.billing.subscription.api.user.DefaultSubscriptionBaseBundle;
 import com.ning.billing.subscription.api.user.DefaultSubscriptionStatusDryRun;
-import com.ning.billing.subscription.api.SubscriptionBase;
+import com.ning.billing.subscription.api.user.EntitlementStatusDryRun;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
-import com.ning.billing.subscription.api.user.SubscriptionBaseTransition;
-import com.ning.billing.subscription.api.user.SubscriptionBuilder;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
-import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
-import com.ning.billing.subscription.api.user.SubscriptionStatusDryRun;
-import com.ning.billing.subscription.api.user.SubscriptionStatusDryRun.DryRunChangeReason;
+import com.ning.billing.subscription.api.user.SubscriptionBaseTransition;
 import com.ning.billing.subscription.api.user.SubscriptionBaseTransitionData;
+import com.ning.billing.subscription.api.user.SubscriptionBuilder;
+import com.ning.billing.subscription.api.user.EntitlementStatusDryRun.DryRunChangeReason;
 import com.ning.billing.subscription.engine.addon.AddonUtils;
 import com.ning.billing.subscription.engine.dao.SubscriptionDao;
 import com.ning.billing.subscription.exceptions.SubscriptionBaseError;
@@ -99,7 +98,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
             final PlanPhase phase = plan.getAllPhases()[0];
             if (phase == null) {
                 throw new SubscriptionBaseError(String.format("No initial PlanPhase for Product %s, term %s and set %s does not exist in the catalog",
-                                                          spec.getProductName(), spec.getBillingPeriod().toString(), realPriceList));
+                                                              spec.getProductName(), spec.getBillingPeriod().toString(), realPriceList));
             }
 
             final SubscriptionBaseBundle bundle = dao.getSubscriptionBundleFromId(bundleId, context);
@@ -142,7 +141,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
                     break;
                 default:
                     throw new SubscriptionBaseError(String.format("Can't create subscription of type %s",
-                                                              plan.getProduct().getCategory().toString()));
+                                                                  plan.getProduct().getCategory().toString()));
             }
 
             return apiService.createPlan(new SubscriptionBuilder()
@@ -178,15 +177,21 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
     }
 
     @Override
+    public List<SubscriptionBaseBundle> getBundlesForKey(final String bundleKey, final InternalTenantContext context) {
+        final List<SubscriptionBaseBundle> result = dao.getSubscriptionBundlesForKey(bundleKey, context);
+        return result;
+    }
+
+    @Override
     public List<SubscriptionBase> getSubscriptionsForBundle(UUID bundleId,
-                                                        InternalTenantContext context) {
+                                                            InternalTenantContext context) {
         final List<SubscriptionBase> internalSubscriptions = dao.getSubscriptions(bundleId, context);
         return createSubscriptionsForApiUse(internalSubscriptions);
     }
 
     @Override
     public SubscriptionBase getBaseSubscription(UUID bundleId,
-                                            InternalTenantContext context) throws SubscriptionBaseApiException {
+                                                InternalTenantContext context) throws SubscriptionBaseApiException {
         final SubscriptionBase result = dao.getBaseSubscription(bundleId, context);
         if (result == null) {
             throw new SubscriptionBaseApiException(ErrorCode.SUB_GET_NO_SUCH_BASE_SUBSCRIPTION, bundleId);
@@ -197,7 +202,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
     @Override
 
     public SubscriptionBase getSubscriptionFromId(UUID id,
-                                              InternalTenantContext context) throws SubscriptionBaseApiException {
+                                                  InternalTenantContext context) throws SubscriptionBaseApiException {
         final SubscriptionBase result = dao.getSubscriptionFromId(id, context);
         if (result == null) {
             throw new SubscriptionBaseApiException(ErrorCode.SUB_INVALID_SUBSCRIPTION_ID, id);
@@ -260,7 +265,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
     }
 
     @Override
-    public List<SubscriptionStatusDryRun> getDryRunChangePlanStatus(final UUID subscriptionId, @Nullable final String baseProductName, final DateTime requestedDate, final InternalTenantContext context) throws SubscriptionBaseApiException {
+    public List<EntitlementStatusDryRun> getDryRunChangePlanStatus(final UUID subscriptionId, @Nullable final String baseProductName, final DateTime requestedDate, final InternalTenantContext context) throws SubscriptionBaseApiException {
         final SubscriptionBase subscription = dao.getSubscriptionFromId(subscriptionId, context);
         if (subscription == null) {
             throw new SubscriptionBaseApiException(ErrorCode.SUB_INVALID_SUBSCRIPTION_ID, subscriptionId);
@@ -269,7 +274,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
             throw new SubscriptionBaseApiException(ErrorCode.SUB_CHANGE_DRY_RUN_NOT_BP);
         }
 
-        final List<SubscriptionStatusDryRun> result = new LinkedList<SubscriptionStatusDryRun>();
+        final List<EntitlementStatusDryRun> result = new LinkedList<EntitlementStatusDryRun>();
 
         final List<SubscriptionBase> bundleSubscriptions = dao.getSubscriptions(subscription.getBundleId(), context);
         for (final SubscriptionBase cur : bundleSubscriptions) {
@@ -291,7 +296,7 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
             } else {
                 reason = DryRunChangeReason.AO_NOT_AVAILABLE_IN_NEW_PLAN;
             }
-            final SubscriptionStatusDryRun status = new DefaultSubscriptionStatusDryRun(cur.getId(),
+            final EntitlementStatusDryRun status = new DefaultSubscriptionStatusDryRun(cur.getId(),
                                                                                         cur.getCurrentPlan().getProduct().getName(),
                                                                                         cur.getCurrentPhase().getPhaseType(),
                                                                                         cur.getCurrentPlan().getBillingPeriod(),
diff --git a/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionStatusDryRun.java b/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionStatusDryRun.java
index 3e7b654..b4e314c 100644
--- a/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionStatusDryRun.java
+++ b/subscription/src/main/java/com/ning/billing/subscription/api/user/DefaultSubscriptionStatusDryRun.java
@@ -20,7 +20,7 @@ import java.util.UUID;
 import com.ning.billing.catalog.api.BillingPeriod;
 import com.ning.billing.catalog.api.PhaseType;
 
-public class DefaultSubscriptionStatusDryRun implements SubscriptionStatusDryRun {
+public class DefaultSubscriptionStatusDryRun implements EntitlementStatusDryRun {
 
     private final UUID id;
     private final String productName;
diff --git a/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiAddOn.java b/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiAddOn.java
index 8d9e597..52a13d2 100644
--- a/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiAddOn.java
+++ b/subscription/src/test/java/com/ning/billing/subscription/api/user/TestUserApiAddOn.java
@@ -43,7 +43,7 @@ import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.entitlement.api.Entitlement.EntitlementState;
 import com.ning.billing.subscription.SubscriptionTestSuiteWithEmbeddedDB;
 import com.ning.billing.subscription.api.SubscriptionBaseTransitionType;
-import com.ning.billing.subscription.api.user.SubscriptionStatusDryRun.DryRunChangeReason;
+import com.ning.billing.subscription.api.user.EntitlementStatusDryRun.DryRunChangeReason;
 
 public class TestUserApiAddOn extends SubscriptionTestSuiteWithEmbeddedDB {
 
@@ -333,7 +333,7 @@ public class TestUserApiAddOn extends SubscriptionTestSuiteWithEmbeddedDB {
             final BillingPeriod newBaseTerm = BillingPeriod.MONTHLY;
             final String newBasePriceList = PriceListSet.DEFAULT_PRICELIST_NAME;
 
-            final List<SubscriptionStatusDryRun> aoStatus = subscriptionInternalApi.getDryRunChangePlanStatus(baseSubscription.getId(),
+            final List<EntitlementStatusDryRun> aoStatus = subscriptionInternalApi.getDryRunChangePlanStatus(baseSubscription.getId(),
                                                                                                      newBaseProduct, now, internalCallContext);
             assertEquals(aoStatus.size(), 1);
             assertEquals(aoStatus.get(0).getId(), aoSubscription.getId());
@@ -396,7 +396,7 @@ public class TestUserApiAddOn extends SubscriptionTestSuiteWithEmbeddedDB {
             final BillingPeriod newBaseTerm = BillingPeriod.MONTHLY;
             final String newBasePriceList = PriceListSet.DEFAULT_PRICELIST_NAME;
 
-            final List<SubscriptionStatusDryRun> aoStatus = subscriptionInternalApi.getDryRunChangePlanStatus(baseSubscription.getId(),
+            final List<EntitlementStatusDryRun> aoStatus = subscriptionInternalApi.getDryRunChangePlanStatus(baseSubscription.getId(),
                                                                                                      newBaseProduct, now, internalCallContext);
             assertEquals(aoStatus.size(), 1);
             assertEquals(aoStatus.get(0).getId(), aoSubscription.getId());
diff --git a/util/src/main/java/com/ning/billing/util/svcapi/subscription/SubscriptionBaseInternalApi.java b/util/src/main/java/com/ning/billing/util/svcapi/subscription/SubscriptionBaseInternalApi.java
index 0303156..1d01828 100644
--- a/util/src/main/java/com/ning/billing/util/svcapi/subscription/SubscriptionBaseInternalApi.java
+++ b/util/src/main/java/com/ning/billing/util/svcapi/subscription/SubscriptionBaseInternalApi.java
@@ -25,8 +25,8 @@ import org.joda.time.DateTime;
 
 import com.ning.billing.catalog.api.PlanPhaseSpecifier;
 import com.ning.billing.subscription.api.SubscriptionBase;
+import com.ning.billing.subscription.api.user.EntitlementStatusDryRun;
 import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
-import com.ning.billing.subscription.api.user.SubscriptionStatusDryRun;
 import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.util.callcontext.InternalCallContext;
 import com.ning.billing.util.callcontext.InternalTenantContext;
@@ -47,6 +47,8 @@ public interface SubscriptionBaseInternalApi {
 
     public List<SubscriptionBaseBundle> getBundlesForAccount(final UUID accountId, final InternalTenantContext context);
 
+    public List<SubscriptionBaseBundle> getBundlesForKey(final String bundleKey, final InternalTenantContext context);
+
     public List<SubscriptionBase> getSubscriptionsForBundle(final UUID bundleId, final InternalTenantContext context);
 
     public SubscriptionBase getBaseSubscription(final UUID bundleId, final InternalTenantContext context) throws SubscriptionBaseApiException;
@@ -65,6 +67,6 @@ public interface SubscriptionBaseInternalApi {
 
     public DateTime getNextBillingDate(final UUID accountId, final InternalTenantContext context);
 
-    public List<SubscriptionStatusDryRun> getDryRunChangePlanStatus(final UUID subscriptionId, @Nullable final String baseProductName,
+    public List<EntitlementStatusDryRun> getDryRunChangePlanStatus(final UUID subscriptionId, @Nullable final String baseProductName,
                                                                     final DateTime requestedDate, final InternalTenantContext context) throws SubscriptionBaseApiException;
 }