killbill-memoizeit

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 a164bef..c1474e4 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
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -29,6 +29,8 @@ import org.killbill.billing.account.api.AccountEmail;
 import org.killbill.billing.account.api.AccountUserApi;
 import org.killbill.billing.account.api.DefaultAccount;
 import org.killbill.billing.account.api.DefaultAccountEmail;
+import org.killbill.billing.account.api.ImmutableAccountData;
+import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.killbill.billing.account.dao.AccountDao;
 import org.killbill.billing.account.dao.AccountEmailModelDao;
 import org.killbill.billing.account.dao.AccountModelDao;
@@ -51,15 +53,18 @@ import static org.killbill.billing.util.entity.dao.DefaultPaginationHelper.getEn
 
 public class DefaultAccountUserApi extends DefaultAccountApiBase implements AccountUserApi {
 
+    private final ImmutableAccountInternalApi immutableAccountInternalApi;
     private final InternalCallContextFactory internalCallContextFactory;
     private final AccountDao accountDao;
 
     @Inject
-    public DefaultAccountUserApi(final AccountDao accountDao,
+    public DefaultAccountUserApi(final ImmutableAccountInternalApi immutableAccountInternalApi,
+                                 final AccountDao accountDao,
                                  final NonEntityDao nonEntityDao,
                                  final CacheControllerDispatcher cacheControllerDispatcher,
                                  final InternalCallContextFactory internalCallContextFactory) {
         super(accountDao, nonEntityDao, cacheControllerDispatcher);
+        this.immutableAccountInternalApi = immutableAccountInternalApi;
         this.internalCallContextFactory = internalCallContextFactory;
         this.accountDao = accountDao;
     }
@@ -89,7 +94,10 @@ public class DefaultAccountUserApi extends DefaultAccountApiBase implements Acco
 
         if (data.getParentAccountId() != null) {
             // verify that parent account exists if parentAccountId is not null
-            getAccountById(data.getParentAccountId(), internalContext);
+            final ImmutableAccountData immutableAccountData = immutableAccountInternalApi.getImmutableAccountDataById(data.getParentAccountId(), internalContext);
+            if (immutableAccountData == null) {
+                throw new AccountApiException(ErrorCode.ACCOUNT_DOES_NOT_EXIST_FOR_ID, data.getParentAccountId());
+            }
         }
 
         final AccountModelDao account = new AccountModelDao(data);

diff --git a/account/src/main/java/org/killbill/billing/account/dao/AccountModelDao.java b/account/src/main/java/org/killbill/billing/account/dao/AccountModelDao.java
index 91fe1ce..30e39f7 100644
--- a/account/src/main/java/org/killbill/billing/account/dao/AccountModelDao.java
+++ b/account/src/main/java/org/killbill/billing/account/dao/AccountModelDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -31,12 +31,13 @@ import org.killbill.billing.util.UUIDs;
 import org.killbill.billing.util.dao.TableName;
 import org.killbill.billing.util.entity.dao.EntityModelDao;
 import org.killbill.billing.util.entity.dao.EntityModelDaoBase;
+import org.killbill.billing.util.entity.dao.TimeZoneAwareEntity;
 
 import com.google.common.base.MoreObjects;
 
 import static org.killbill.billing.account.api.DefaultMutableAccountData.DEFAULT_BILLING_CYCLE_DAY_LOCAL;
 
-public class AccountModelDao extends EntityModelDaoBase implements EntityModelDao<Account> {
+public class AccountModelDao extends EntityModelDaoBase implements TimeZoneAwareEntity, EntityModelDao<Account> {
 
     private String externalKey;
     private String email;
@@ -205,6 +206,7 @@ public class AccountModelDao extends EntityModelDaoBase implements EntityModelDa
         this.paymentMethodId = paymentMethodId;
     }
 
+    @Override
     public DateTimeZone getTimeZone() {
         return timeZone;
     }

diff --git a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
index 6ca2fa3..171e157 100644
--- a/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
+++ b/account/src/main/java/org/killbill/billing/account/dao/DefaultAccountDao.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -91,7 +91,7 @@ public class DefaultAccountDao extends EntityDaoBase<AccountModelDao, Account, A
 
         final Long recordId = entitySqlDaoWrapperFactory.become(AccountSqlDao.class).getRecordId(savedAccount.getId().toString(), context);
         // We need to re-hydrate the callcontext with the account record id
-        final InternalCallContext rehydratedContext = internalCallContextFactory.createInternalCallContext(recordId, context);
+        final InternalCallContext rehydratedContext = internalCallContextFactory.createInternalCallContext(savedAccount, recordId, context);
         final AccountCreationInternalEvent creationEvent = new DefaultAccountCreationEvent(new DefaultAccountData(savedAccount), savedAccount.getId(),
                                                                                            rehydratedContext.getAccountRecordId(), rehydratedContext.getTenantRecordId(), rehydratedContext.getUserToken());
         try {

diff --git a/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java b/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
index 807e5e1..447b18d 100644
--- a/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
+++ b/account/src/test/java/org/killbill/billing/account/AccountTestSuiteNoDB.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * 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:
  *
@@ -16,6 +18,7 @@
 
 package org.killbill.billing.account;
 
+import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
@@ -43,6 +46,8 @@ public abstract class AccountTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
     @Inject
     protected AccountDao accountDao;
     @Inject
+    protected ImmutableAccountInternalApi immutableAccountInternalApi;
+    @Inject
     protected AccountUserApi accountUserApi;
     @Inject
     protected AuditDao auditDao;

diff --git a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
index 9fe5b8c..b5dbac0 100644
--- a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
+++ b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApi.java
@@ -350,12 +350,10 @@ public class TestDefaultAccountUserApi extends AccountTestSuiteWithEmbeddedDB {
     @Test(groups = "slow", description = "Test Account create Child with a non existing Parent",
             expectedExceptions = AccountApiException.class, expectedExceptionsMessageRegExp = "Account does not exist for id .*")
     public void testCreateChildAccountWithInvalidParent() throws Exception {
-
         final AccountModelDao childAccountModel = createTestAccount();
         childAccountModel.setParentAccountId(UUID.randomUUID());
         final AccountData childAccountData = new DefaultAccount(childAccountModel);
         final Account childAccount = accountUserApi.createAccount(childAccountData, callContext);
-
     }
 
     @Test(groups = "slow", description = "Test un- and re-parenting")

diff --git a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApiWithMocks.java b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApiWithMocks.java
index 793d689..44e370c 100644
--- a/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApiWithMocks.java
+++ b/account/src/test/java/org/killbill/billing/account/api/user/TestDefaultAccountUserApiWithMocks.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -53,7 +53,7 @@ public class TestDefaultAccountUserApiWithMocks extends AccountTestSuiteNoDB {
     @BeforeMethod(groups = "fast")
     public void setUp() throws Exception {
         accountDao = new MockAccountDao(Mockito.mock(PersistentBus.class), clock);
-        accountUserApi = new DefaultAccountUserApi(accountDao, nonEntityDao, controllerDispatcher, internalFactory);
+        accountUserApi = new DefaultAccountUserApi(immutableAccountInternalApi, accountDao, nonEntityDao, controllerDispatcher, internalFactory);
     }
 
     @Test(groups = "fast", description = "Test Account create API")

diff --git a/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotifier.java b/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotifier.java
index 6d083d1..51ddb3a 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotifier.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotifier.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * 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:
  *
@@ -20,14 +22,14 @@ import java.util.UUID;
 import java.util.concurrent.Callable;
 
 import org.joda.time.DateTime;
-import org.testng.Assert;
-import org.testng.annotations.Test;
-
-import org.killbill.billing.subscription.api.SubscriptionBase;
+import org.killbill.billing.ObjectType;
+import org.killbill.billing.account.api.Account;
 import org.killbill.billing.invoice.InvoiceTestSuiteWithEmbeddedDB;
 import org.killbill.billing.invoice.api.DefaultInvoiceService;
+import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.notificationq.api.NotificationQueue;
-import org.killbill.clock.ClockMock;
+import org.testng.Assert;
+import org.testng.annotations.Test;
 
 import static com.jayway.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.MINUTES;
@@ -36,19 +38,19 @@ public class TestNextBillingDateNotifier extends InvoiceTestSuiteWithEmbeddedDB 
 
     @Test(groups = "slow")
     public void testInvoiceNotifier() throws Exception {
+        final Account account = invoiceUtil.createAccount(callContext);
+        final Long accountRecordId = nonEntityDao.retrieveAccountRecordIdFromObject(account.getId(), ObjectType.ACCOUNT, null);
 
         final SubscriptionBase subscription = invoiceUtil.createSubscription();
         final UUID subscriptionId = subscription.getId();
         final DateTime now = clock.getUTCNow();
 
-
         final NotificationQueue nextBillingQueue = notificationQueueService.getNotificationQueue(DefaultInvoiceService.INVOICE_SERVICE_NAME, DefaultNextBillingDateNotifier.NEXT_BILLING_DATE_NOTIFIER_QUEUE);
 
-
-        nextBillingQueue.recordFutureNotification(now, new NextBillingDateNotificationKey(subscriptionId, now, Boolean.FALSE), internalCallContext.getUserToken(), internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
+        nextBillingQueue.recordFutureNotification(now, new NextBillingDateNotificationKey(subscriptionId, now, Boolean.FALSE), internalCallContext.getUserToken(), accountRecordId, internalCallContext.getTenantRecordId());
 
         // Move time in the future after the notification effectiveDate
-        ((ClockMock) clock).setDeltaFromReality(3000);
+        clock.setDeltaFromReality(3000);
 
         await().atMost(1, MINUTES).until(new Callable<Boolean>() {
             @Override

diff --git a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
index 7708eb6..06ee314 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -52,8 +52,6 @@ import org.killbill.billing.catalog.api.PlanPhaseSpecifier;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.entitlement.api.SubscriptionEventType;
 import org.killbill.billing.entity.EntityPersistenceException;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.api.DryRunArguments;
 import org.killbill.billing.invoice.api.DryRunType;
 import org.killbill.billing.invoice.api.Invoice;
@@ -95,7 +93,6 @@ import org.testng.Assert;
 import com.google.common.base.Function;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 
 public class TestInvoiceHelper {
 

diff --git a/overdue/src/test/java/org/killbill/billing/overdue/notification/TestOverdueCheckNotifier.java b/overdue/src/test/java/org/killbill/billing/overdue/notification/TestOverdueCheckNotifier.java
index ee470c5..9fee915 100644
--- a/overdue/src/test/java/org/killbill/billing/overdue/notification/TestOverdueCheckNotifier.java
+++ b/overdue/src/test/java/org/killbill/billing/overdue/notification/TestOverdueCheckNotifier.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * 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:
  *
@@ -20,18 +22,17 @@ import java.util.UUID;
 import java.util.concurrent.Callable;
 
 import org.joda.time.DateTime;
-import org.mockito.Mockito;
-import org.testng.Assert;
-import org.testng.annotations.AfterMethod;
-import org.testng.annotations.BeforeMethod;
-import org.testng.annotations.Test;
-
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.callcontext.InternalCallContext;
+import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.overdue.OverdueTestSuiteWithEmbeddedDB;
 import org.killbill.billing.overdue.listener.OverdueDispatcher;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
-import org.killbill.billing.callcontext.InternalTenantContext;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
 
 import static com.jayway.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -41,9 +42,7 @@ public class TestOverdueCheckNotifier extends OverdueTestSuiteWithEmbeddedDB {
     private OverdueDispatcherMock mockDispatcher;
     private OverdueNotifier notifierForMock;
 
-
-
-    private static final class  OverdueDispatcherMock extends OverdueDispatcher {
+    private static final class OverdueDispatcherMock extends OverdueDispatcher {
 
         int eventCount = 0;
         UUID latestAccountId = null;
@@ -73,9 +72,6 @@ public class TestOverdueCheckNotifier extends OverdueTestSuiteWithEmbeddedDB {
         //super.beforeMethod();
         // We override the parent method on purpose, because we want to register a different OverdueCheckNotifier
 
-        final Account account = Mockito.mock(Account.class);
-        Mockito.when(accountApi.getAccountById(Mockito.<UUID>any(), Mockito.<InternalTenantContext>any())).thenReturn(account);
-
         mockDispatcher = new OverdueDispatcherMock(internalCallContextFactory);
         notifierForMock = new OverdueCheckNotifier(notificationQueueService, overdueProperties, internalCallContextFactory, mockDispatcher);
 
@@ -95,6 +91,7 @@ public class TestOverdueCheckNotifier extends OverdueTestSuiteWithEmbeddedDB {
         final UUID accountId = new UUID(0L, 1L);
         final Account account = Mockito.mock(Account.class);
         Mockito.when(account.getId()).thenReturn(accountId);
+        Mockito.when(accountApi.getImmutableAccountDataByRecordId(Mockito.<UUID>eq(internalCallContext.getAccountRecordId()), Mockito.<InternalTenantContext>any())).thenReturn(account);
         final DateTime now = clock.getUTCNow();
         final DateTime readyTime = now.plusMillis(2000);
 

diff --git a/overdue/src/test/java/org/killbill/billing/overdue/OverdueTestSuiteWithEmbeddedDB.java b/overdue/src/test/java/org/killbill/billing/overdue/OverdueTestSuiteWithEmbeddedDB.java
index f2b590d..b74d629 100644
--- a/overdue/src/test/java/org/killbill/billing/overdue/OverdueTestSuiteWithEmbeddedDB.java
+++ b/overdue/src/test/java/org/killbill/billing/overdue/OverdueTestSuiteWithEmbeddedDB.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014 Groupon, Inc
- * Copyright 2014 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -21,7 +21,7 @@ package org.killbill.billing.overdue;
 import javax.inject.Named;
 
 import org.killbill.billing.GuicyKillbillTestSuiteWithEmbeddedDB;
-import org.killbill.billing.account.api.AccountInternalApi;
+import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.killbill.billing.invoice.api.InvoiceInternalApi;
 import org.killbill.billing.junction.BlockingInternalApi;
 import org.killbill.billing.lifecycle.api.BusService;
@@ -52,7 +52,7 @@ import com.google.inject.Injector;
 public abstract class OverdueTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuiteWithEmbeddedDB {
 
     @Inject
-    protected AccountInternalApi accountApi;
+    protected ImmutableAccountInternalApi accountApi;
     @Inject
     protected BillingStateCalculator calculatorBundle;
     @Inject

diff --git a/subscription/src/test/java/org/killbill/billing/subscription/api/user/TestUserApiError.java b/subscription/src/test/java/org/killbill/billing/subscription/api/user/TestUserApiError.java
index cd663dd..968376a 100644
--- a/subscription/src/test/java/org/killbill/billing/subscription/api/user/TestUserApiError.java
+++ b/subscription/src/test/java/org/killbill/billing/subscription/api/user/TestUserApiError.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -77,6 +77,9 @@ public class TestUserApiError extends SubscriptionTestSuiteNoDB {
     @Test(groups = "fast")
     public void testCreateSubscriptionAddOnNotAvailable() throws SubscriptionBaseApiException {
         final SubscriptionBaseBundle aoBundle = subscriptionInternalApi.createBundleForAccount(bundle.getAccountId(), "myAOBundle", internalCallContext);
+        mockNonEntityDao.addTenantRecordIdMapping(aoBundle.getId(), internalCallContext);
+        mockNonEntityDao.addAccountRecordIdMapping(aoBundle.getId(), internalCallContext);
+
         testUtil.createSubscriptionWithBundle(aoBundle.getId(), "Pistol", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         tCreateSubscriptionInternal(aoBundle.getId(), "Telescopic-Scope", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, ErrorCode.SUB_CREATE_AO_NOT_AVAILABLE);
     }
@@ -84,6 +87,9 @@ public class TestUserApiError extends SubscriptionTestSuiteNoDB {
     @Test(groups = "fast")
     public void testCreateSubscriptionAddOnIncluded() throws SubscriptionBaseApiException {
         final SubscriptionBaseBundle aoBundle = subscriptionInternalApi.createBundleForAccount(bundle.getAccountId(), "myAOBundle", internalCallContext);
+        mockNonEntityDao.addTenantRecordIdMapping(aoBundle.getId(), internalCallContext);
+        mockNonEntityDao.addAccountRecordIdMapping(aoBundle.getId(), internalCallContext);
+
         testUtil.createSubscriptionWithBundle(aoBundle.getId(), "Assault-Rifle", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         tCreateSubscriptionInternal(aoBundle.getId(), "Telescopic-Scope", BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, ErrorCode.SUB_CREATE_AO_ALREADY_INCLUDED);
     }

diff --git a/subscription/src/test/java/org/killbill/billing/subscription/SubscriptionTestSuiteNoDB.java b/subscription/src/test/java/org/killbill/billing/subscription/SubscriptionTestSuiteNoDB.java
index 0ace9f9..b90587c 100644
--- a/subscription/src/test/java/org/killbill/billing/subscription/SubscriptionTestSuiteNoDB.java
+++ b/subscription/src/test/java/org/killbill/billing/subscription/SubscriptionTestSuiteNoDB.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -24,7 +24,10 @@ import javax.inject.Inject;
 
 import org.killbill.billing.GuicyKillbillTestSuiteNoDB;
 import org.killbill.billing.account.api.AccountData;
+import org.killbill.billing.account.api.ImmutableAccountData;
+import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.killbill.billing.api.TestApiListener;
+import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Catalog;
 import org.killbill.billing.catalog.api.CatalogService;
 import org.killbill.billing.dao.MockNonEntityDao;
@@ -62,6 +65,8 @@ public class SubscriptionTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
     protected static final Logger log = LoggerFactory.getLogger(SubscriptionTestSuiteNoDB.class);
 
     @Inject
+    protected ImmutableAccountInternalApi immutableAccountInternalApi;
+    @Inject
     protected SubscriptionBaseService subscriptionBaseService;
     @Inject
     protected SubscriptionBaseInternalApi subscriptionInternalApi;
@@ -132,7 +137,13 @@ public class SubscriptionTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
         this.accountData = subscriptionTestInitializer.initAccountData();
         final UUID accountId = UUIDs.randomUUID();
         mockNonEntityDao.addTenantRecordIdMapping(accountId, internalCallContext);
+
+        final ImmutableAccountData immutableAccountData = Mockito.mock(ImmutableAccountData.class);
+        Mockito.when(immutableAccountInternalApi.getImmutableAccountDataByRecordId(Mockito.<Long>eq(internalCallContext.getAccountRecordId()), Mockito.<InternalTenantContext>any())).thenReturn(immutableAccountData);
+
         this.bundle = subscriptionTestInitializer.initBundle(accountId, subscriptionInternalApi, internalCallContext);
+        mockNonEntityDao.addTenantRecordIdMapping(bundle.getId(), internalCallContext);
+        mockNonEntityDao.addAccountRecordIdMapping(bundle.getId(), internalCallContext);
     }
 
     @AfterMethod(groups = "fast")

diff --git a/util/src/main/java/org/killbill/billing/util/account/AccountDateTimeUtils.java b/util/src/main/java/org/killbill/billing/util/account/AccountDateTimeUtils.java
index 6ab4322..58f9157 100644
--- a/util/src/main/java/org/killbill/billing/util/account/AccountDateTimeUtils.java
+++ b/util/src/main/java/org/killbill/billing/util/account/AccountDateTimeUtils.java
@@ -1,6 +1,6 @@
 /*
- * Copyright 2016 Groupon, Inc
- * Copyright 2016 The Billing Project, LLC
+ * Copyright 2016-2017 Groupon, Inc
+ * Copyright 2016-2017 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
@@ -20,11 +20,20 @@ package org.killbill.billing.util.account;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.killbill.billing.account.api.Account;
+import org.killbill.billing.util.entity.Entity;
+import org.killbill.billing.util.entity.dao.TimeZoneAwareEntity;
 
 public abstract class AccountDateTimeUtils {
 
+    public static DateTimeZone getFixedOffsetTimeZone(final TimeZoneAwareEntity account) {
+        return getFixedOffsetTimeZone(account.getTimeZone(), account);
+    }
+
     public static DateTimeZone getFixedOffsetTimeZone(final Account account) {
-        final DateTimeZone referenceDateTimeZone = account.getTimeZone();
+        return getFixedOffsetTimeZone(account.getTimeZone(), account);
+    }
+
+    public static DateTimeZone getFixedOffsetTimeZone(final DateTimeZone referenceDateTimeZone, final Entity account) {
         final DateTime referenceDateTime = getReferenceDateTime(account);
 
         // Check if DST was in effect at the reference date time
@@ -36,7 +45,7 @@ public abstract class AccountDateTimeUtils {
         }
     }
 
-    public static DateTime getReferenceDateTime(final Account account) {
+    public static DateTime getReferenceDateTime(final Entity account) {
         return account.getCreatedDate();
     }
 }

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 ffee721..a42fe24 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
@@ -31,14 +31,17 @@ import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
+import org.killbill.billing.util.account.AccountDateTimeUtils;
 import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.cache.CacheController;
 import org.killbill.billing.util.cache.CacheControllerDispatcher;
 import org.killbill.billing.util.dao.NonEntityDao;
+import org.killbill.billing.util.entity.dao.TimeZoneAwareEntity;
 import org.killbill.clock.Clock;
 import org.slf4j.MDC;
 
 import com.google.common.base.MoreObjects;
+import com.google.common.base.Preconditions;
 
 // Internal contexts almost always expect accountRecordId and tenantRecordId to be populated
 public class InternalCallContextFactory {
@@ -153,7 +156,10 @@ public class InternalCallContextFactory {
         if (accountRecordId == null) {
             return new InternalTenantContext(tenantRecordId);
         } else {
-            return new InternalTenantContext(tenantRecordId, accountRecordId, getFixedOffsetTimeZone(accountRecordId, tenantRecordId), getReferenceTime(accountRecordId, tenantRecordId));
+            final ImmutableAccountData immutableAccountData = getImmutableAccountData(accountRecordId, tenantRecordId);
+            final DateTimeZone fixedOffsetTimeZone = immutableAccountData.getFixedOffsetTimeZone();
+            final DateTime referenceTime = immutableAccountData.getReferenceTime();
+            return new InternalTenantContext(tenantRecordId, accountRecordId, fixedOffsetTimeZone, referenceTime);
         }
     }
 
@@ -237,8 +243,23 @@ public class InternalCallContextFactory {
 
     // Used when we need to re-hydrate the callcontext with the account_record_id (when creating the account)
     public InternalCallContext createInternalCallContext(final Long accountRecordId, final InternalCallContext context) {
-        final DateTimeZone fixedOffsetTimeZone = getFixedOffsetTimeZone(accountRecordId, context.getTenantRecordId());
-        final DateTime referenceTime = getReferenceTime(accountRecordId, context.getTenantRecordId());
+        final ImmutableAccountData immutableAccountData = getImmutableAccountData(accountRecordId, context.getTenantRecordId());
+        final DateTimeZone fixedOffsetTimeZone = immutableAccountData.getFixedOffsetTimeZone();
+        final DateTime referenceTime = immutableAccountData.getReferenceTime();
+        populateMDCContext(accountRecordId, context.getTenantRecordId());
+        return new InternalCallContext(context, accountRecordId, fixedOffsetTimeZone, referenceTime, clock.getUTCNow());
+    }
+
+    // Used during the account creation transaction (account not visible outside of the transaction yet)
+    public InternalCallContext createInternalCallContext(final TimeZoneAwareEntity accountModelDao, final Long accountRecordId, final InternalCallContext context) {
+        // See DefaultImmutableAccountData implementation
+        final DateTimeZone fixedOffsetTimeZone = AccountDateTimeUtils.getFixedOffsetTimeZone(accountModelDao);
+        final DateTime referenceTime = AccountDateTimeUtils.getReferenceDateTime(accountModelDao);
+        populateMDCContext(accountRecordId, context.getTenantRecordId());
+        return new InternalCallContext(context, accountRecordId, fixedOffsetTimeZone, referenceTime, clock.getUTCNow());
+    }
+
+    public InternalCallContext createInternalCallContext(final DateTimeZone fixedOffsetTimeZone, final DateTime referenceTime, final Long accountRecordId, final InternalCallContext context) {
         populateMDCContext(accountRecordId, context.getTenantRecordId());
         return new InternalCallContext(context, accountRecordId, fixedOffsetTimeZone, referenceTime, clock.getUTCNow());
     }
@@ -257,8 +278,18 @@ public class InternalCallContextFactory {
                                                           final CallOrigin callOrigin, final UserType userType, @Nullable final UUID userToken,
                                                           @Nullable final String reasonCode, @Nullable final String comment) {
         final Long nonNulTenantRecordId = MoreObjects.firstNonNull(tenantRecordId, INTERNAL_TENANT_RECORD_ID);
-        final DateTimeZone fixedOffsetTimeZone = getFixedOffsetTimeZone(accountRecordId, tenantRecordId);
-        final DateTime referenceTime = getReferenceTime(accountRecordId, tenantRecordId);
+
+        final DateTimeZone fixedOffsetTimeZone;
+        final DateTime referenceTime;
+        if (accountRecordId == null) {
+            // TENANT_CONFIG_CHANGE event for instance
+            fixedOffsetTimeZone = null;
+            referenceTime = null;
+        } else {
+            final ImmutableAccountData immutableAccountData = getImmutableAccountData(accountRecordId, nonNulTenantRecordId);
+            fixedOffsetTimeZone = immutableAccountData.getFixedOffsetTimeZone();
+            referenceTime = immutableAccountData.getReferenceTime();
+        }
 
         populateMDCContext(accountRecordId, nonNulTenantRecordId);
 
@@ -276,34 +307,15 @@ public class InternalCallContextFactory {
                                        clock.getUTCNow());
     }
 
-    private DateTimeZone getFixedOffsetTimeZone(@Nullable final Long accountRecordId, final Long tenantRecordId) {
-        if (accountRecordId == null || accountInternalApi == null) {
-            return null;
-        }
-
-        populateMDCContext(accountRecordId, tenantRecordId);
-
-        final ImmutableAccountData immutableAccountData = getImmutableAccountData(accountRecordId, tenantRecordId);
-        // Will be null while creating the account
-        return immutableAccountData == null ? null : immutableAccountData.getFixedOffsetTimeZone();
-    }
-
-    private DateTime getReferenceTime(@Nullable final Long accountRecordId, final Long tenantRecordId) {
-        if (accountRecordId == null || accountInternalApi == null) {
-            return null;
-        }
-
-        final ImmutableAccountData immutableAccountData = getImmutableAccountData(accountRecordId, tenantRecordId);
-        // Will be null while creating the account
-        return immutableAccountData == null ? null : immutableAccountData.getReferenceTime();
-    }
-
-    private ImmutableAccountData getImmutableAccountData(@Nullable final Long accountRecordId, final Long tenantRecordId) {
+    private ImmutableAccountData getImmutableAccountData(final Long accountRecordId, final Long tenantRecordId) {
+        Preconditions.checkNotNull(accountRecordId, "Missing accountRecordId");
         final InternalTenantContext tmp = new InternalTenantContext(tenantRecordId, accountRecordId, null, null);
         try {
-            return accountInternalApi.getImmutableAccountDataByRecordId(accountRecordId, tmp);
+            final ImmutableAccountData immutableAccountData = accountInternalApi.getImmutableAccountDataByRecordId(accountRecordId, tmp);
+            Preconditions.checkNotNull(immutableAccountData, "Unable to retrieve immutableAccountData");
+            return immutableAccountData;
         } catch (final AccountApiException e) {
-            return null;
+            throw new RuntimeException(e);
         }
     }
 

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
index 280d8ac..2ddeb1c 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
@@ -21,6 +21,7 @@ package org.killbill.billing.util.entity.dao;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
+import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.sql.PreparedStatement;
 import java.sql.SQLException;
@@ -293,9 +294,14 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
             int foundIndexForEntityModelDao = -1;
             for (int i = 0; i < types.length; i++) {
                 final Class clz = ((Class) types[i]);
-                if (EntityModelDao.class.getName().equals(((Class) ((java.lang.reflect.ParameterizedType) clz.getGenericInterfaces()[0]).getRawType()).getName())) {
-                    foundIndexForEntityModelDao = i;
-                    break;
+                final Type[] genericInterfaces = clz.getGenericInterfaces();
+                for (final Type genericInterface : genericInterfaces) {
+                    if (genericInterface instanceof ParameterizedType) {
+                        if (EntityModelDao.class.getName().equals(((Class) ((ParameterizedType) genericInterface).getRawType()).getName())) {
+                            foundIndexForEntityModelDao = i;
+                            break;
+                        }
+                    }
                 }
             }
 
@@ -389,7 +395,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
                 }
 
                 // Make sure to re-hydrate the object (especially needed for create calls)
-                insertAudits(tableName, entityRecordId, historyRecordId, changeType, context);
+                insertAudits(tableName, reHydratedEntity, entityRecordId, historyRecordId, changeType, context);
                 return null;
             }
         });
@@ -463,14 +469,16 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
         return nonEntityDao.retrieveLastHistoryRecordIdFromTransaction(entityRecordId, entityModelDao.getHistoryTableName(), transactional);
     }
 
-    private void insertAudits(final TableName tableName, final Long entityRecordId, final Long historyRecordId, final ChangeType changeType, final InternalCallContext contextMaybeWithoutAccountRecordId) {
+    private void insertAudits(final TableName tableName, final M entityModelDao, final Long entityRecordId, final Long historyRecordId, final ChangeType changeType, final InternalCallContext contextMaybeWithoutAccountRecordId) {
         final TableName destinationTableName = MoreObjects.firstNonNull(tableName.getHistoryTableName(), tableName);
         final EntityAudit audit = new EntityAudit(destinationTableName, historyRecordId, changeType, clock.getUTCNow());
 
         final InternalCallContext context;
         // Populate the account record id when creating the account record
         if (TableName.ACCOUNT.equals(tableName) && ChangeType.INSERT.equals(changeType)) {
-            context = internalCallContextFactory.createInternalCallContext(entityRecordId, contextMaybeWithoutAccountRecordId);
+            // AccountModelDao in practice
+            final TimeZoneAwareEntity accountModelDao = (TimeZoneAwareEntity) entityModelDao;
+            context = internalCallContextFactory.createInternalCallContext(accountModelDao, entityRecordId, contextMaybeWithoutAccountRecordId);
         } else {
             context = contextMaybeWithoutAccountRecordId;
         }

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/TimeZoneAwareEntity.java b/util/src/main/java/org/killbill/billing/util/entity/dao/TimeZoneAwareEntity.java
new file mode 100644
index 0000000..feff107
--- /dev/null
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/TimeZoneAwareEntity.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2017 Groupon, Inc
+ * Copyright 2017 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.util.entity.dao;
+
+import org.joda.time.DateTimeZone;
+import org.killbill.billing.util.entity.Entity;
+
+public interface TimeZoneAwareEntity extends Entity {
+
+    public DateTimeZone getTimeZone();
+}

diff --git a/util/src/test/java/org/killbill/billing/util/callcontext/TestInternalCallContextFactory.java b/util/src/test/java/org/killbill/billing/util/callcontext/TestInternalCallContextFactory.java
index 9a5c01e..40e940e 100644
--- a/util/src/test/java/org/killbill/billing/util/callcontext/TestInternalCallContextFactory.java
+++ b/util/src/test/java/org/killbill/billing/util/callcontext/TestInternalCallContextFactory.java
@@ -1,7 +1,9 @@
 /*
- * Copyright 2010-2012 Ning, Inc.
+ * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * 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:
  *
@@ -19,6 +21,9 @@ package org.killbill.billing.util.callcontext;
 import java.util.Date;
 import java.util.UUID;
 
+import org.killbill.billing.account.api.ImmutableAccountData;
+import org.killbill.billing.callcontext.InternalTenantContext;
+import org.mockito.Mockito;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.testng.Assert;
@@ -36,6 +41,9 @@ public class TestInternalCallContextFactory extends UtilTestSuiteWithEmbeddedDB 
         final UUID invoiceId = UUID.randomUUID();
         final Long accountRecordId = 19384012L;
 
+        final ImmutableAccountData immutableAccountData = Mockito.mock(ImmutableAccountData.class);
+        Mockito.when(immutableAccountInternalApi.getImmutableAccountDataByRecordId(Mockito.<Long>eq(accountRecordId), Mockito.<InternalTenantContext>any())).thenReturn(immutableAccountData);
+
         dbi.withHandle(new HandleCallback<Void>() {
             @Override
             public Void withHandle(final Handle handle) throws Exception {
@@ -71,6 +79,9 @@ public class TestInternalCallContextFactory extends UtilTestSuiteWithEmbeddedDB 
         final UUID accountId = UUID.randomUUID();
         final Long accountRecordId = 19384012L;
 
+        final ImmutableAccountData immutableAccountData = Mockito.mock(ImmutableAccountData.class);
+        Mockito.when(immutableAccountInternalApi.getImmutableAccountDataByRecordId(Mockito.<Long>eq(accountRecordId), Mockito.<InternalTenantContext>any())).thenReturn(immutableAccountData);
+
         dbi.withHandle(new HandleCallback<Void>() {
             @Override
             public Void withHandle(final Handle handle) throws Exception {

diff --git a/util/src/test/java/org/killbill/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java b/util/src/test/java/org/killbill/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java
index 2dccd4c..ffe76e3 100644
--- a/util/src/test/java/org/killbill/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java
+++ b/util/src/test/java/org/killbill/billing/util/customfield/api/TestDefaultCustomFieldUserApi.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -24,11 +24,14 @@ import java.util.Map;
 import java.util.UUID;
 
 import org.killbill.billing.ObjectType;
+import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.api.TestApiListener.NextEvent;
+import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.util.UtilTestSuiteWithEmbeddedDB;
 import org.killbill.billing.util.customfield.CustomField;
 import org.killbill.billing.util.customfield.StringCustomField;
 import org.killbill.billing.util.entity.Pagination;
+import org.mockito.Mockito;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.testng.Assert;
@@ -43,6 +46,9 @@ public class TestDefaultCustomFieldUserApi extends UtilTestSuiteWithEmbeddedDB {
         final UUID accountId = UUID.randomUUID();
         final Long accountRecordId = 19384012L;
 
+        final ImmutableAccountData immutableAccountData = Mockito.mock(ImmutableAccountData.class);
+        Mockito.when(immutableAccountInternalApi.getImmutableAccountDataByRecordId(Mockito.<Long>eq(accountRecordId), Mockito.<InternalTenantContext>any())).thenReturn(immutableAccountData);
+
         dbi.withHandle(new HandleCallback<Void>() {
             @Override
             public Void withHandle(final Handle handle) throws Exception {

diff --git a/util/src/test/java/org/killbill/billing/util/tag/api/TestDefaultTagUserApi.java b/util/src/test/java/org/killbill/billing/util/tag/api/TestDefaultTagUserApi.java
index 978f5bb..bbdfb0f 100644
--- a/util/src/test/java/org/killbill/billing/util/tag/api/TestDefaultTagUserApi.java
+++ b/util/src/test/java/org/killbill/billing/util/tag/api/TestDefaultTagUserApi.java
@@ -1,6 +1,6 @@
 /*
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -23,11 +23,14 @@ import java.util.Map;
 import java.util.UUID;
 
 import org.killbill.billing.ObjectType;
+import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.api.TestApiListener.NextEvent;
+import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.util.UtilTestSuiteWithEmbeddedDB;
 import org.killbill.billing.util.entity.Pagination;
 import org.killbill.billing.util.tag.ControlTagType;
 import org.killbill.billing.util.tag.Tag;
+import org.mockito.Mockito;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.testng.Assert;
@@ -42,6 +45,9 @@ public class TestDefaultTagUserApi extends UtilTestSuiteWithEmbeddedDB {
         final UUID accountId = UUID.randomUUID();
         final Long accountRecordId = 19384012L;
 
+        final ImmutableAccountData immutableAccountData = Mockito.mock(ImmutableAccountData.class);
+        Mockito.when(immutableAccountInternalApi.getImmutableAccountDataByRecordId(Mockito.<Long>eq(accountRecordId), Mockito.<InternalTenantContext>any())).thenReturn(immutableAccountData);
+
         dbi.withHandle(new HandleCallback<Void>() {
             @Override
             public Void withHandle(final Handle handle) throws Exception {

diff --git a/util/src/test/java/org/killbill/billing/util/UtilTestSuiteWithEmbeddedDB.java b/util/src/test/java/org/killbill/billing/util/UtilTestSuiteWithEmbeddedDB.java
index 17ce189..503a229 100644
--- a/util/src/test/java/org/killbill/billing/util/UtilTestSuiteWithEmbeddedDB.java
+++ b/util/src/test/java/org/killbill/billing/util/UtilTestSuiteWithEmbeddedDB.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2014 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 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
@@ -21,6 +21,7 @@ package org.killbill.billing.util;
 import javax.inject.Inject;
 
 import org.killbill.billing.GuicyKillbillTestSuiteWithEmbeddedDB;
+import org.killbill.billing.account.api.ImmutableAccountInternalApi;
 import org.killbill.billing.api.TestApiListener;
 import org.killbill.billing.security.api.SecurityApi;
 import org.killbill.billing.util.audit.dao.AuditDao;
@@ -66,6 +67,8 @@ public abstract class UtilTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuite
     @Inject
     protected InternalCallContextFactory internalCallContextFactory;
     @Inject
+    protected ImmutableAccountInternalApi immutableAccountInternalApi;
+    @Inject
     protected DefaultTagUserApi tagUserApi;
     @Inject
     protected DefaultCustomFieldUserApi customFieldUserApi;