killbill-memoizeit

diff --git a/account/src/test/java/org/killbill/billing/account/dao/MockAccountDao.java b/account/src/test/java/org/killbill/billing/account/dao/MockAccountDao.java
index 9d097c4..077ca8d 100644
--- a/account/src/test/java/org/killbill/billing/account/dao/MockAccountDao.java
+++ b/account/src/test/java/org/killbill/billing/account/dao/MockAccountDao.java
@@ -37,6 +37,8 @@ import org.killbill.billing.account.api.user.DefaultAccountCreationEvent.Default
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.events.AccountChangeInternalEvent;
+import org.killbill.billing.util.api.AuditLevel;
+import org.killbill.billing.util.audit.AuditLogWithHistory;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
 import org.killbill.billing.util.entity.DefaultPagination;
 import org.killbill.billing.util.entity.Pagination;
@@ -186,4 +188,9 @@ public class MockAccountDao extends MockEntityDaoBase<AccountModelDao, Account, 
             }
         }));
     }
+
+    @Override
+    public List<AuditLogWithHistory> getAuditLogsWithHistoryForId(final UUID accountId, final AuditLevel auditLevel, final InternalTenantContext context) throws AccountApiException {
+        throw new UnsupportedOperationException();
+    }
 }

diff --git a/account/src/test/java/org/killbill/billing/account/dao/TestAccountDao.java b/account/src/test/java/org/killbill/billing/account/dao/TestAccountDao.java
index 3db923f..9d4408f 100644
--- a/account/src/test/java/org/killbill/billing/account/dao/TestAccountDao.java
+++ b/account/src/test/java/org/killbill/billing/account/dao/TestAccountDao.java
@@ -26,7 +26,6 @@ import org.joda.time.DateTimeZone;
 import org.killbill.billing.ErrorCode;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.account.AccountTestSuiteWithEmbeddedDB;
-import org.killbill.billing.account.api.Account;
 import org.killbill.billing.account.api.AccountApiException;
 import org.killbill.billing.account.api.AccountData;
 import org.killbill.billing.account.api.AccountEmail;
@@ -40,10 +39,10 @@ import org.killbill.billing.util.api.CustomFieldApiException;
 import org.killbill.billing.util.api.TagApiException;
 import org.killbill.billing.util.api.TagDefinitionApiException;
 import org.killbill.billing.util.audit.AuditLog;
+import org.killbill.billing.util.audit.AuditLogWithHistory;
 import org.killbill.billing.util.audit.ChangeType;
 import org.killbill.billing.util.audit.DefaultAccountAuditLogs;
 import org.killbill.billing.util.customfield.dao.CustomFieldModelDao;
-import org.killbill.billing.util.dao.EntityHistoryModelDao;
 import org.killbill.billing.util.dao.TableName;
 import org.killbill.billing.util.entity.Pagination;
 import org.killbill.billing.util.tag.DescriptiveTag;
@@ -214,17 +213,18 @@ public class TestAccountDao extends AccountTestSuiteWithEmbeddedDB {
         final AccountModelDao account = createTestAccount();
         accountDao.create(account, internalCallContext);
         final AccountModelDao createdAccount = accountDao.getAccountByKey(account.getExternalKey(), internalCallContext);
-
-        final List<EntityHistoryModelDao<AccountModelDao, Account>> history1 = getAccountHistory(createdAccount.getRecordId());
-        Assert.assertEquals(history1.size(), 1);
-        Assert.assertEquals(history1.get(0).getChangeType(), ChangeType.INSERT);
-        Assert.assertEquals(history1.get(0).getEntity().getAccountRecordId(), createdAccount.getRecordId());
-        Assert.assertEquals(history1.get(0).getEntity().getTenantRecordId(), createdAccount.getTenantRecordId());
-        Assert.assertEquals(history1.get(0).getEntity().getExternalKey(), createdAccount.getExternalKey());
-        Assert.assertEquals(history1.get(0).getEntity().getMigrated(), createdAccount.getMigrated());
-        Assert.assertEquals(history1.get(0).getEntity().getIsNotifiedForInvoices(), createdAccount.getIsNotifiedForInvoices());
-        Assert.assertEquals(history1.get(0).getEntity().getTimeZone(), createdAccount.getTimeZone());
-        Assert.assertEquals(history1.get(0).getEntity().getLocale(), createdAccount.getLocale());
+        List<AuditLogWithHistory> auditLogsWithHistory = accountDao.getAuditLogsWithHistoryForId(account.getId(), AuditLevel.FULL, internalCallContext);
+
+        final AccountModelDao history1 = (AccountModelDao) auditLogsWithHistory.get(0).getEntity();
+        Assert.assertEquals(auditLogsWithHistory.size(), 1);
+        Assert.assertEquals(auditLogsWithHistory.get(0).getChangeType(), ChangeType.INSERT);
+        Assert.assertEquals(history1.getAccountRecordId(), createdAccount.getRecordId());
+        Assert.assertEquals(history1.getTenantRecordId(), createdAccount.getTenantRecordId());
+        Assert.assertEquals(history1.getExternalKey(), createdAccount.getExternalKey());
+        Assert.assertEquals(history1.getMigrated(), createdAccount.getMigrated());
+        Assert.assertEquals(history1.getIsNotifiedForInvoices(), createdAccount.getIsNotifiedForInvoices());
+        Assert.assertEquals(history1.getTimeZone(), createdAccount.getTimeZone());
+        Assert.assertEquals(history1.getLocale(), createdAccount.getLocale());
 
         final AccountData accountData = new MockAccountBuilder(new DefaultAccount(account)).migrated(false)
                                                                                            .isNotifiedForInvoices(false)
@@ -236,18 +236,19 @@ public class TestAccountDao extends AccountTestSuiteWithEmbeddedDB {
 
         final AccountModelDao retrievedAccount = accountDao.getAccountByKey(account.getExternalKey(), internalCallContext);
         checkAccountsEqual(retrievedAccount, updatedAccount);
-
-        final List<EntityHistoryModelDao<AccountModelDao, Account>> history2 = getAccountHistory(createdAccount.getRecordId());
-        Assert.assertEquals(history2.size(), 2);
-        Assert.assertEquals(history2.get(0).getChangeType(), ChangeType.INSERT);
-        Assert.assertEquals(history2.get(1).getChangeType(), ChangeType.UPDATE);
-        Assert.assertEquals(history2.get(1).getEntity().getAccountRecordId(), retrievedAccount.getRecordId());
-        Assert.assertEquals(history2.get(1).getEntity().getTenantRecordId(), retrievedAccount.getTenantRecordId());
-        Assert.assertEquals(history2.get(1).getEntity().getExternalKey(), retrievedAccount.getExternalKey());
-        Assert.assertEquals(history2.get(1).getEntity().getMigrated(), retrievedAccount.getMigrated());
-        Assert.assertEquals(history2.get(1).getEntity().getIsNotifiedForInvoices(), retrievedAccount.getIsNotifiedForInvoices());
-        Assert.assertEquals(history2.get(1).getEntity().getTimeZone(), retrievedAccount.getTimeZone());
-        Assert.assertEquals(history2.get(1).getEntity().getLocale(), retrievedAccount.getLocale());
+        auditLogsWithHistory = accountDao.getAuditLogsWithHistoryForId(retrievedAccount.getId(), AuditLevel.FULL, internalCallContext);
+
+        final AccountModelDao history2 = (AccountModelDao) auditLogsWithHistory.get(1).getEntity();
+        Assert.assertEquals(auditLogsWithHistory.size(), 2);
+        Assert.assertEquals(auditLogsWithHistory.get(0).getChangeType(), ChangeType.INSERT);
+        Assert.assertEquals(auditLogsWithHistory.get(1).getChangeType(), ChangeType.UPDATE);
+        Assert.assertEquals(history2.getAccountRecordId(), retrievedAccount.getRecordId());
+        Assert.assertEquals(history2.getTenantRecordId(), retrievedAccount.getTenantRecordId());
+        Assert.assertEquals(history2.getExternalKey(), retrievedAccount.getExternalKey());
+        Assert.assertEquals(history2.getMigrated(), retrievedAccount.getMigrated());
+        Assert.assertEquals(history2.getIsNotifiedForInvoices(), retrievedAccount.getIsNotifiedForInvoices());
+        Assert.assertEquals(history2.getTimeZone(), retrievedAccount.getTimeZone());
+        Assert.assertEquals(history2.getLocale(), retrievedAccount.getLocale());
 
         final AccountData accountData2 = new MockAccountBuilder(new DefaultAccount(updatedAccount)).isNotifiedForInvoices(true)
                                                                                                    .locale("en_US")
@@ -257,19 +258,20 @@ public class TestAccountDao extends AccountTestSuiteWithEmbeddedDB {
 
         final AccountModelDao retrievedAccount2 = accountDao.getAccountByKey(account.getExternalKey(), internalCallContext);
         checkAccountsEqual(retrievedAccount2, updatedAccount2);
-
-        final List<EntityHistoryModelDao<AccountModelDao, Account>> history3 = getAccountHistory(createdAccount.getRecordId());
-        Assert.assertEquals(history3.size(), 3);
-        Assert.assertEquals(history3.get(0).getChangeType(), ChangeType.INSERT);
-        Assert.assertEquals(history3.get(1).getChangeType(), ChangeType.UPDATE);
-        Assert.assertEquals(history3.get(2).getChangeType(), ChangeType.UPDATE);
-        Assert.assertEquals(history3.get(2).getEntity().getAccountRecordId(), retrievedAccount2.getRecordId());
-        Assert.assertEquals(history3.get(2).getEntity().getTenantRecordId(), retrievedAccount2.getTenantRecordId());
-        Assert.assertEquals(history3.get(2).getEntity().getExternalKey(), retrievedAccount2.getExternalKey());
-        Assert.assertEquals(history3.get(2).getEntity().getMigrated(), retrievedAccount2.getMigrated());
-        Assert.assertEquals(history3.get(2).getEntity().getIsNotifiedForInvoices(), retrievedAccount2.getIsNotifiedForInvoices());
-        Assert.assertEquals(history3.get(2).getEntity().getTimeZone(), retrievedAccount2.getTimeZone());
-        Assert.assertEquals(history3.get(2).getEntity().getLocale(), retrievedAccount2.getLocale());
+        auditLogsWithHistory = accountDao.getAuditLogsWithHistoryForId(retrievedAccount2.getId(), AuditLevel.FULL, internalCallContext);
+
+        final AccountModelDao history3 = (AccountModelDao) auditLogsWithHistory.get(2).getEntity();
+        Assert.assertEquals(auditLogsWithHistory.size(), 3);
+        Assert.assertEquals(auditLogsWithHistory.get(0).getChangeType(), ChangeType.INSERT);
+        Assert.assertEquals(auditLogsWithHistory.get(1).getChangeType(), ChangeType.UPDATE);
+        Assert.assertEquals(auditLogsWithHistory.get(2).getChangeType(), ChangeType.UPDATE);
+        Assert.assertEquals(history3.getAccountRecordId(), retrievedAccount2.getRecordId());
+        Assert.assertEquals(history3.getTenantRecordId(), retrievedAccount2.getTenantRecordId());
+        Assert.assertEquals(history3.getExternalKey(), retrievedAccount2.getExternalKey());
+        Assert.assertEquals(history3.getMigrated(), retrievedAccount2.getMigrated());
+        Assert.assertEquals(history3.getIsNotifiedForInvoices(), retrievedAccount2.getIsNotifiedForInvoices());
+        Assert.assertEquals(history3.getTimeZone(), retrievedAccount2.getTimeZone());
+        Assert.assertEquals(history3.getLocale(), retrievedAccount2.getLocale());
     }
 
     @Test(groups = "slow", description = "Test Account DAO: payment method update")
@@ -399,10 +401,4 @@ public class TestAccountDao extends AccountTestSuiteWithEmbeddedDB {
         Assert.assertEquals(auditLogsForAccountEmail2.size(), 1);
         Assert.assertEquals(auditLogsForAccountEmail2.get(0).getChangeType(), ChangeType.INSERT);
     }
-
-    private List<EntityHistoryModelDao<AccountModelDao, Account>> getAccountHistory(final Long accountRecordId) {
-        // See https://github.com/killbill/killbill/issues/335
-        final AccountSqlDao accountSqlDao = dbi.onDemand(AccountSqlDao.class);
-        return accountSqlDao.getHistoryForTargetRecordId(accountRecordId, internalCallContext);
-    }
 }

diff --git a/util/src/test/java/org/killbill/billing/mock/api/MockAccountUserApi.java b/util/src/test/java/org/killbill/billing/mock/api/MockAccountUserApi.java
index d6798fc..cbbf33a 100644
--- a/util/src/test/java/org/killbill/billing/mock/api/MockAccountUserApi.java
+++ b/util/src/test/java/org/killbill/billing/mock/api/MockAccountUserApi.java
@@ -30,6 +30,8 @@ import org.killbill.billing.account.api.AccountEmail;
 import org.killbill.billing.account.api.AccountUserApi;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.mock.MockAccountBuilder;
+import org.killbill.billing.util.api.AuditLevel;
+import org.killbill.billing.util.audit.AuditLogWithHistory;
 import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.billing.util.entity.DefaultPagination;
@@ -187,4 +189,9 @@ public class MockAccountUserApi implements AccountUserApi {
     public List<Account> getChildrenAccounts(final UUID uuid, final TenantContext tenantContext) throws AccountApiException {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public List<AuditLogWithHistory> getAuditLogsWithHistoryForId(final UUID uuid, final AuditLevel auditLevel, final TenantContext tenantContext) throws AccountApiException {
+        throw new UnsupportedOperationException();
+    }
 }

diff --git a/util/src/test/java/org/killbill/billing/util/audit/dao/MockAuditDao.java b/util/src/test/java/org/killbill/billing/util/audit/dao/MockAuditDao.java
index 1aa816d..86cbfc9 100644
--- a/util/src/test/java/org/killbill/billing/util/audit/dao/MockAuditDao.java
+++ b/util/src/test/java/org/killbill/billing/util/audit/dao/MockAuditDao.java
@@ -25,9 +25,11 @@ import java.util.UUID;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.util.api.AuditLevel;
 import org.killbill.billing.util.audit.AuditLog;
+import org.killbill.billing.util.audit.AuditLogWithHistory;
 import org.killbill.billing.util.audit.DefaultAccountAuditLogs;
 import org.killbill.billing.util.audit.DefaultAccountAuditLogsForObjectType;
 import org.killbill.billing.util.dao.TableName;
+import org.killbill.billing.util.entity.dao.EntitySqlDao;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
@@ -81,4 +83,9 @@ public class MockAuditDao implements AuditDao {
             return allAuditLogs;
         }
     }
+
+    @Override
+    public List<AuditLogWithHistory> getAuditLogsWithHistoryForId(final EntitySqlDao sqlDao, final TableName tableName, final UUID objectId, final AuditLevel auditLevel, final InternalTenantContext context) {
+        throw new UnsupportedOperationException();
+    }
 }