killbill-memoizeit

context: populate accountRecordId when possible InternalCallContextFactory#createInternalTenantContext(TenantContext)

2/10/2016 8:53:02 AM

Details

diff --git a/account/src/main/java/org/killbill/billing/account/api/user/DefaultAccountUserApi.java b/account/src/main/java/org/killbill/billing/account/api/user/DefaultAccountUserApi.java
index b137cc0..bef0089 100644
--- a/account/src/main/java/org/killbill/billing/account/api/user/DefaultAccountUserApi.java
+++ b/account/src/main/java/org/killbill/billing/account/api/user/DefaultAccountUserApi.java
@@ -72,7 +72,7 @@ public class DefaultAccountUserApi extends DefaultAccountApiBase implements Acco
 
     @Override
     public Account getAccountById(final UUID id, final TenantContext context) throws AccountApiException {
-        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(context);
+        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(id, context);
         return getAccountById(id, internalTenantContext);
     }
 
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
index 3d43c6b..6526e06 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
@@ -92,7 +92,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
     protected EventsStream eventsStream;
 
 
-    public DefaultEntitlement(final UUID entitlementId, final EventsStreamBuilder eventsStreamBuilder,
+    public DefaultEntitlement(final UUID accountId, final UUID entitlementId, final EventsStreamBuilder eventsStreamBuilder,
                               final EntitlementApi entitlementApi, final EntitlementPluginExecution pluginExecution, final BlockingStateDao blockingStateDao,
                               final SubscriptionBaseInternalApi subscriptionInternalApi, final BlockingChecker checker,
                               final NotificationQueueService notificationQueueService, final EntitlementUtils entitlementUtils,
@@ -100,7 +100,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
                               final InternalCallContextFactory internalCallContextFactory, final TenantContext tenantContext) throws EntitlementApiException {
         this(eventsStreamBuilder.buildForEntitlement(entitlementId, tenantContext), eventsStreamBuilder,
              entitlementApi, pluginExecution, blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
-             entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory.createInternalTenantContext(tenantContext), internalCallContextFactory);
+             entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory.createInternalTenantContext(accountId, tenantContext), internalCallContextFactory);
     }
 
     public DefaultEntitlement(final EventsStream eventsStream, final EventsStreamBuilder eventsStreamBuilder,
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
index 6e6cd45..1678364 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
@@ -146,7 +146,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                     final EntitlementSpecifier specifier = getFirstEntitlementSpecifier(updatedPluginContext.getEntitlementSpecifiers());
                     final SubscriptionBase subscription = subscriptionBaseInternalApi.createSubscription(bundle.getId(), specifier.getPlanPhaseSpecifier(), specifier.getOverrides(), requestedDate, contextWithValidAccountRecordId);
 
-                    return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
+                    return new DefaultEntitlement(accountId, subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
                 } catch (final SubscriptionBaseApiException e) {
@@ -209,7 +209,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                     final DateTime requestedDate = dateHelper.fromLocalDateAndReferenceTime(updatedPluginContext.getEffectiveDate(), referenceTime, contextWithValidAccountRecordId);
                     final SubscriptionBase subscription = subscriptionBaseInternalApi.createBaseSubscriptionWithAddOns(bundle.getId(), entitlementSpecifiers, requestedDate, contextWithValidAccountRecordId);
 
-                    return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
+                    return new DefaultEntitlement(accountId, subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
 
@@ -260,7 +260,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                     final EntitlementSpecifier specifier = getFirstEntitlementSpecifier(updatedPluginContext.getEntitlementSpecifiers());
                     final SubscriptionBase subscription = subscriptionBaseInternalApi.createSubscription(bundleId, specifier.getPlanPhaseSpecifier(), specifier.getOverrides(), requestedDate, context);
 
-                    return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
+                    return new DefaultEntitlement(eventsStreamForBaseSubscription.getAccountId(), subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
                 } catch (final SubscriptionBaseApiException e) {
@@ -273,7 +273,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
 
     @Override
     public List<EntitlementAOStatusDryRun> getDryRunStatusForChange(final UUID bundleId, final String targetProductName, final LocalDate effectiveDate, final TenantContext context) throws EntitlementApiException {
-        final InternalTenantContext internalContext = internalCallContextFactory.createInternalTenantContext(context);
+        final InternalTenantContext internalContext = internalCallContextFactory.createInternalTenantContext(bundleId, ObjectType.BUNDLE, context);
         try {
             final SubscriptionBaseBundle bundle = subscriptionBaseInternalApi.getBundleFromId(bundleId, internalContext);
             final SubscriptionBase baseSubscription = subscriptionBaseInternalApi.getBaseSubscription(bundleId, internalContext);
@@ -294,7 +294,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
 
     @Override
     public List<Entitlement> getAllEntitlementsForBundle(final UUID bundleId, final TenantContext tenantContext) throws EntitlementApiException {
-        final InternalTenantContext internalContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
+        final InternalTenantContext internalContext = internalCallContextFactory.createInternalTenantContext(bundleId, ObjectType.BUNDLE, tenantContext);
         final UUID accountId;
         try {
             accountId = subscriptionBaseInternalApi.getBundleFromId(bundleId, internalContext).getAccountId();
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultSubscriptionApi.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultSubscriptionApi.java
index b4fe99c..95d15ac 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultSubscriptionApi.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultSubscriptionApi.java
@@ -121,10 +121,6 @@ public class DefaultSubscriptionApi implements SubscriptionApi {
 
     @Override
     public SubscriptionBundle getSubscriptionBundle(final UUID bundleId, final TenantContext tenantContext) throws SubscriptionApiException {
-
-        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
-
-
         final UUID accountId = internalCallContextFactory.getAccountId(bundleId, ObjectType.BUNDLE, tenantContext);
 
         final Optional<SubscriptionBundle> bundleOptional = Iterables.<SubscriptionBundle>tryFind(getSubscriptionBundlesForAccount(accountId, tenantContext),
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
index b090a03..b2cd20a 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
@@ -100,7 +100,7 @@ public class EventsStreamBuilder {
     public EventsStream buildForBaseSubscription(final UUID bundleId, final TenantContext tenantContext) throws EntitlementApiException {
         final SubscriptionBase baseSubscription;
         try {
-            final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(tenantContext);
+            final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(bundleId, ObjectType.BUNDLE, tenantContext);
             baseSubscription = subscriptionInternalApi.getBaseSubscription(bundleId, internalTenantContext);
         } catch (SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java b/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
index 9e1cd9d..7e7406c 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
@@ -127,7 +127,7 @@ public class DefaultInvoiceUserApi implements InvoiceUserApi {
 
     @Override
     public Invoice getInvoiceByPayment(final UUID paymentId, final TenantContext context) throws InvoiceApiException {
-        final InternalTenantContext tenantContext = internalCallContextFactory.createInternalTenantContext(context);
+        final InternalTenantContext tenantContext = internalCallContextFactory.createInternalTenantContext(paymentId, ObjectType.PAYMENT, context);
         final UUID invoiceId = dao.getInvoiceIdByPaymentId(paymentId, tenantContext);
         if (invoiceId == null) {
             throw new InvoiceApiException(ErrorCode.INVOICE_NOT_FOUND, paymentId);
diff --git a/jaxrs/src/main/java/org/killbill/billing/jaxrs/mappers/IllegalStateExceptionMapper.java b/jaxrs/src/main/java/org/killbill/billing/jaxrs/mappers/IllegalStateExceptionMapper.java
new file mode 100644
index 0000000..1b9167c
--- /dev/null
+++ b/jaxrs/src/main/java/org/killbill/billing/jaxrs/mappers/IllegalStateExceptionMapper.java
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2016 Groupon, Inc
+ * Copyright 2016 The Billing Project, LLC
+ *
+ * The Billing Project 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 org.killbill.billing.jaxrs.mappers;
+
+import javax.inject.Singleton;
+import javax.ws.rs.core.Context;
+import javax.ws.rs.core.Response;
+import javax.ws.rs.core.UriInfo;
+import javax.ws.rs.ext.ExceptionMapper;
+import javax.ws.rs.ext.Provider;
+
+import org.killbill.billing.util.callcontext.InternalCallContextFactory.ObjectDoesNotExist;
+
+@Singleton
+@Provider
+public class IllegalStateExceptionMapper extends ExceptionMapperBase implements ExceptionMapper<IllegalStateException> {
+
+    private final UriInfo uriInfo;
+
+    public IllegalStateExceptionMapper(@Context final UriInfo uriInfo) {
+        this.uriInfo = uriInfo;
+    }
+
+    @Override
+    public Response toResponse(final IllegalStateException exception) {
+        if (exception instanceof ObjectDoesNotExist) {
+            // Likely object for wrong tenant
+            return buildNotFoundResponse(exception, uriInfo);
+        } else {
+            return buildInternalErrorResponse(exception, uriInfo);
+        }
+    }
+}
diff --git a/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java b/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
index a2701e4..3e6aad6 100644
--- a/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
+++ b/overdue/src/main/java/org/killbill/billing/overdue/api/DefaultOverdueInternalApi.java
@@ -84,8 +84,8 @@ public class DefaultOverdueInternalApi implements OverdueInternalApi {
     @Override
     public OverdueState getOverdueStateFor(final ImmutableAccountData overdueable, final TenantContext context) throws OverdueException {
         try {
-            final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(context);
-            final BlockingState blockingStateForService = accessApi.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalCallContextFactory.createInternalTenantContext(context));
+            final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(overdueable.getId(), context);
+            final BlockingState blockingStateForService = accessApi.getBlockingStateForService(overdueable.getId(), BlockingStateType.ACCOUNT, OverdueService.OVERDUE_SERVICE_NAME, internalTenantContext);
             final String stateName = blockingStateForService != null ? blockingStateForService.getStateName() : OverdueWrapper.CLEAR_STATE_NAME;
             final OverdueConfig overdueConfig = overdueConfigCache.getOverdueConfig(internalTenantContext);
             final OverdueStateSet states = ((DefaultOverdueConfig) overdueConfig).getOverdueStatesAccount();
@@ -99,7 +99,7 @@ public class DefaultOverdueInternalApi implements OverdueInternalApi {
     public BillingState getBillingStateFor(final ImmutableAccountData overdueable, final TenantContext context) throws OverdueException {
         log.debug("Billing state of of {} requested", overdueable.getId());
 
-        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(context);
+        final InternalTenantContext internalTenantContext = internalCallContextFactory.createInternalTenantContext(overdueable.getId(), context);
         final OverdueWrapper wrapper = factory.createOverdueWrapperFor(overdueable, internalTenantContext);
         return wrapper.billingState(internalTenantContext);
     }
diff --git a/payment/src/main/java/org/killbill/billing/payment/api/DefaultPaymentApi.java b/payment/src/main/java/org/killbill/billing/payment/api/DefaultPaymentApi.java
index 3da498d..ccb946e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/api/DefaultPaymentApi.java
+++ b/payment/src/main/java/org/killbill/billing/payment/api/DefaultPaymentApi.java
@@ -26,6 +26,7 @@ import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.killbill.billing.ErrorCode;
+import org.killbill.billing.ObjectType;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.catalog.api.Currency;
@@ -392,7 +393,7 @@ public class DefaultPaymentApi extends DefaultApiBase implements PaymentApi {
 
     @Override
     public Payment getPayment(final UUID paymentId, final boolean withPluginInfo, final Iterable<PluginProperty> properties, final TenantContext context) throws PaymentApiException {
-        final Payment payment = paymentProcessor.getPayment(paymentId, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(context));
+        final Payment payment = paymentProcessor.getPayment(paymentId, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(paymentId, ObjectType.PAYMENT, context));
         if (payment == null) {
             throw new PaymentApiException(ErrorCode.PAYMENT_NO_SUCH_PAYMENT, paymentId);
         }
@@ -431,13 +432,13 @@ public class DefaultPaymentApi extends DefaultApiBase implements PaymentApi {
     @Override
     public List<PaymentMethod> getAccountPaymentMethods(final UUID accountId, final boolean withPluginInfo, final Iterable<PluginProperty> properties, final TenantContext context)
             throws PaymentApiException {
-        return paymentMethodProcessor.getPaymentMethods(accountId, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(context));
+        return paymentMethodProcessor.getPaymentMethods(accountId, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(accountId, context));
     }
 
     @Override
     public PaymentMethod getPaymentMethodById(final UUID paymentMethodId, final boolean includedDeleted, final boolean withPluginInfo, final Iterable<PluginProperty> properties, final TenantContext context)
             throws PaymentApiException {
-        return paymentMethodProcessor.getPaymentMethodById(paymentMethodId, includedDeleted, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(context));
+        return paymentMethodProcessor.getPaymentMethodById(paymentMethodId, includedDeleted, withPluginInfo, properties, context, internalCallContextFactory.createInternalTenantContext(paymentMethodId, ObjectType.PAYMENT_METHOD, context));
     }
 
     @Override
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestBundle.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestBundle.java
index e9bc2be..fbc8a31 100644
--- a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestBundle.java
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestBundle.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -74,12 +74,7 @@ public class TestBundle extends TestJaxrsBase {
     public void testBundleNonExistent() throws Exception {
         final Account accountJson = createAccount();
 
-        try {
-            killBillClient.getBundle(UUID.randomUUID());
-            Assert.fail();
-        } catch (final KillBillClientException e) {
-            Assert.assertEquals(e.getBillingException().getClassName(), "org.killbill.billing.util.callcontext.InternalCallContextFactory$ObjectDoesNotExist");
-        }
+        Assert.assertNull(killBillClient.getBundle(UUID.randomUUID()));
         Assert.assertTrue(killBillClient.getAccountBundles(accountJson.getAccountId(), "98374982743892").isEmpty());
         Assert.assertTrue(killBillClient.getAccountBundles(accountJson.getAccountId()).isEmpty());
     }
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestCredit.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestCredit.java
index 5bc85f0..48c3748 100644
--- a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestCredit.java
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestCredit.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -26,7 +26,6 @@ import org.killbill.billing.client.KillBillClientException;
 import org.killbill.billing.client.model.Account;
 import org.killbill.billing.client.model.Credit;
 import org.killbill.billing.client.model.Invoice;
-import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -88,11 +87,6 @@ public class TestCredit extends TestJaxrsBase {
 
     @Test(groups = "slow", description = "Cannot retrieve a non existing credit")
     public void testCreditDoesNotExist() throws Exception {
-        try {
-            killBillClient.getCredit(UUID.randomUUID());
-            Assert.fail();
-        } catch (final KillBillClientException e) {
-            Assert.assertEquals(e.getBillingException().getClassName(), "org.killbill.billing.util.callcontext.InternalCallContextFactory$ObjectDoesNotExist");
-        }
+        assertNull(killBillClient.getCredit(UUID.randomUUID()));
     }
 }
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestEntitlement.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestEntitlement.java
index c795c16..1d3cb3f 100644
--- a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestEntitlement.java
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestEntitlement.java
@@ -145,26 +145,11 @@ public class TestEntitlement extends TestJaxrsBase {
         subscription.setBillingPeriod(BillingPeriod.ANNUAL);
         subscription.setPriceList(PriceListSet.DEFAULT_PRICELIST_NAME);
 
-        try {
-            killBillClient.updateSubscription(subscription, CALL_COMPLETION_TIMEOUT_SEC, createdBy, reason, comment);
-            Assert.fail();
-        } catch (final KillBillClientException e) {
-            Assert.assertEquals(e.getBillingException().getClassName(), "org.killbill.billing.util.callcontext.InternalCallContextFactory$ObjectDoesNotExist");
-        }
-
-        try {
-            killBillClient.cancelSubscription(subscriptionId, createdBy, reason, comment);
-            Assert.fail();
-        } catch (final KillBillClientException e) {
-            Assert.assertEquals(e.getBillingException().getClassName(), "org.killbill.billing.util.callcontext.InternalCallContextFactory$ObjectDoesNotExist");
-        }
-
-        try {
-            killBillClient.getSubscription(subscriptionId);
-            Assert.fail();
-        } catch (final KillBillClientException e) {
-            Assert.assertEquals(e.getBillingException().getClassName(), "org.killbill.billing.util.callcontext.InternalCallContextFactory$ObjectDoesNotExist");
-        }
+        assertNull(killBillClient.updateSubscription(subscription, CALL_COMPLETION_TIMEOUT_SEC, createdBy, reason, comment));
+
+        killBillClient.cancelSubscription(subscriptionId, createdBy, reason, comment);
+
+        assertNull(killBillClient.getSubscription(subscriptionId));
     }
 
     @Test(groups = "slow", description = "Can override billing policy on change")
diff --git a/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java b/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
index 1feb447..45b0146 100644
--- a/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
+++ b/usage/src/main/java/org/killbill/billing/usage/api/user/DefaultUsageUserApi.java
@@ -64,7 +64,7 @@ public class DefaultUsageUserApi implements UsageUserApi {
 
     @Override
     public RolledUpUsage getUsageForSubscription(final UUID subscriptionId, final String unitType, final LocalDate startDate, final LocalDate endDate, final TenantContext tenantContext) {
-        final List<RolledUpUsageModelDao> usageForSubscription = rolledUpUsageDao.getUsageForSubscription(subscriptionId, startDate, endDate, unitType, internalCallContextFactory.createInternalTenantContext(tenantContext));
+        final List<RolledUpUsageModelDao> usageForSubscription = rolledUpUsageDao.getUsageForSubscription(subscriptionId, startDate, endDate, unitType, internalCallContextFactory.createInternalTenantContext(subscriptionId, ObjectType.SUBSCRIPTION, tenantContext));
         final List<RolledUpUnit> rolledUpAmount = getRolledUpUnits(usageForSubscription);
         return new DefaultRolledUpUsage(subscriptionId, startDate, endDate, rolledUpAmount);
     }
diff --git a/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java b/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
index d4367b5..d5434d1 100644
--- a/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
+++ b/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
@@ -325,7 +325,7 @@ public class InternalCallContextFactory {
         return nonEntityDao.retrieveTenantRecordIdFromObject(objectId, objectType, cacheControllerDispatcher.getCacheController(CacheType.TENANT_RECORD_ID));
     }
 
-    private static final class ObjectDoesNotExist extends IllegalStateException {
+    public static final class ObjectDoesNotExist extends IllegalStateException {
 
         public ObjectDoesNotExist(final String s) {
             super(s);