killbill-memoizeit

diff --git a/account/src/test/java/com/ning/billing/account/dao/TestAccountDao.java b/account/src/test/java/com/ning/billing/account/dao/TestAccountDao.java
index 752a304..c26cb32 100644
--- a/account/src/test/java/com/ning/billing/account/dao/TestAccountDao.java
+++ b/account/src/test/java/com/ning/billing/account/dao/TestAccountDao.java
@@ -33,6 +33,8 @@ import com.ning.billing.account.api.AccountEmail;
 import com.ning.billing.account.api.DefaultAccount;
 import com.ning.billing.account.api.DefaultAccountEmail;
 import com.ning.billing.account.api.MutableAccountData;
+import com.ning.billing.callcontext.InternalCallContext;
+import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.mock.MockAccountBuilder;
 import com.ning.billing.util.api.AuditLevel;
 import com.ning.billing.util.api.CustomFieldApiException;
@@ -40,6 +42,7 @@ import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.api.TagDefinitionApiException;
 import com.ning.billing.util.audit.AuditLog;
 import com.ning.billing.util.audit.ChangeType;
+import com.ning.billing.util.audit.DefaultAccountAuditLogs;
 import com.ning.billing.util.customfield.CustomField;
 import com.ning.billing.util.customfield.StringCustomField;
 import com.ning.billing.util.customfield.dao.CustomFieldModelDao;
@@ -83,6 +86,39 @@ public class TestAccountDao extends AccountTestSuiteWithEmbeddedDB {
         Assert.assertEquals(auditLogsForAccount.get(0).getChangeType(), ChangeType.INSERT);
     }
 
+    @Test(groups = "slow", description = "Test Account: verify audits")
+    public void testAudits() throws AccountApiException {
+        // Special test to verify audits - they are handled a bit differently due to the account record id (see EntitySqlDaoWrapperInvocationHandler#insertAudits)
+        final AccountModelDao account1 = createTestAccount();
+        accountDao.create(account1, internalCallContext);
+        final Long account1RecordId = nonEntityDao.retrieveAccountRecordIdFromObject(account1.getId(), ObjectType.ACCOUNT, null);
+        final InternalCallContext internalCallContext1 = new InternalCallContext(internalCallContext, account1RecordId);
+
+        // Verify audits via account record id
+        final DefaultAccountAuditLogs auditLogsForAccount1ViaAccountRecordId1 = auditDao.getAuditLogsForAccountRecordId(AuditLevel.FULL, internalCallContext1);
+        Assert.assertEquals(auditLogsForAccount1ViaAccountRecordId1.getAuditLogsForAccount().size(), 1);
+        Assert.assertEquals(auditLogsForAccount1ViaAccountRecordId1.getAuditLogsForAccount().get(0).getChangeType(), ChangeType.INSERT);
+
+        // Add an entry in the account_history table to make sure we pick up the right
+        // record id / target record id / account record id in the audit_log table
+        accountDao.updatePaymentMethod(account1.getId(), UUID.randomUUID(), internalCallContext1);
+
+        final AccountModelDao account2 = createTestAccount();
+        accountDao.create(account2, internalCallContext);
+        final Long account2RecordId = nonEntityDao.retrieveAccountRecordIdFromObject(account2.getId(), ObjectType.ACCOUNT, null);
+        final InternalTenantContext internalTenantContext2 = new InternalCallContext(internalCallContext, account2RecordId);
+
+        // Verify audits via account record id
+        final DefaultAccountAuditLogs auditLogsForAccount2ViaAccountRecordId = auditDao.getAuditLogsForAccountRecordId(AuditLevel.FULL, internalTenantContext2);
+        Assert.assertEquals(auditLogsForAccount2ViaAccountRecordId.getAuditLogsForAccount().size(), 1);
+        Assert.assertEquals(auditLogsForAccount2ViaAccountRecordId.getAuditLogsForAccount().get(0).getChangeType(), ChangeType.INSERT);
+
+        final DefaultAccountAuditLogs auditLogsForAccount1ViaAccountRecordId2 = auditDao.getAuditLogsForAccountRecordId(AuditLevel.FULL, internalCallContext1);
+        Assert.assertEquals(auditLogsForAccount1ViaAccountRecordId2.getAuditLogsForAccount().size(), 2);
+        Assert.assertEquals(auditLogsForAccount1ViaAccountRecordId2.getAuditLogsForAccount().get(0).getChangeType(), ChangeType.INSERT);
+        Assert.assertEquals(auditLogsForAccount1ViaAccountRecordId2.getAuditLogsForAccount().get(1).getChangeType(), ChangeType.UPDATE);
+    }
+
     // Simple test to ensure long phone numbers can be stored
     @Test(groups = "slow", description = "Test Account DAO: long numbers")
     public void testLongPhoneNumber() throws AccountApiException {

diff --git a/api/src/main/java/com/ning/billing/callcontext/InternalCallContext.java b/api/src/main/java/com/ning/billing/callcontext/InternalCallContext.java
index fb83e76..ae6e1c9 100644
--- a/api/src/main/java/com/ning/billing/callcontext/InternalCallContext.java
+++ b/api/src/main/java/com/ning/billing/callcontext/InternalCallContext.java
@@ -62,6 +62,12 @@ public class InternalCallContext extends InternalTenantContext {
              callContext.getUpdatedDate());
     }
 
+    public InternalCallContext(final InternalCallContext context, final Long accountRecordId) {
+        this(context.getTenantRecordId(), accountRecordId, context.getUserToken(), context.getCreatedBy(), context.getCallOrigin(),
+             context.getContextUserType(), context.getReasonCode(), context.getComments(), context.getCreatedDate(),
+             context.getUpdatedDate());
+    }
+
     // TODO should not be needed if all services are using internal API
     // Unfortunately not true as some APIs ae hidden in object -- e.g OverdueStateApplicator is doing subscription.cancelEntitlementWithDateOverrideBillingPolicy
     public CallContext toCallContext(final UUID tenantId) {

diff --git a/util/src/main/java/com/ning/billing/util/audit/dao/AuditDao.java b/util/src/main/java/com/ning/billing/util/audit/dao/AuditDao.java
index 294eda3..e11446f 100644
--- a/util/src/main/java/com/ning/billing/util/audit/dao/AuditDao.java
+++ b/util/src/main/java/com/ning/billing/util/audit/dao/AuditDao.java
@@ -28,10 +28,10 @@ import com.ning.billing.util.dao.TableName;
 
 public interface AuditDao {
 
-    // Make sure to consume all or call close() when done for release the connection
+    // Make sure to consume all or call close() when done to release the connection
     public DefaultAccountAuditLogs getAuditLogsForAccountRecordId(AuditLevel auditLevel, InternalTenantContext context);
 
-    // Make sure to consume all or call close() when done for release the connection
+    // Make sure to consume all or call close() when done to release the connection
     public DefaultAccountAuditLogsForObjectType getAuditLogsForAccountRecordId(TableName tableName, AuditLevel auditLevel, InternalTenantContext context);
 
     public List<AuditLog> getAuditLogsForId(TableName tableName, UUID objectId, AuditLevel auditLevel, InternalTenantContext context);

diff --git a/util/src/main/java/com/ning/billing/util/audit/DefaultAccountAuditLogs.java b/util/src/main/java/com/ning/billing/util/audit/DefaultAccountAuditLogs.java
index ef1f691..8515dd5 100644
--- a/util/src/main/java/com/ning/billing/util/audit/DefaultAccountAuditLogs.java
+++ b/util/src/main/java/com/ning/billing/util/audit/DefaultAccountAuditLogs.java
@@ -57,7 +57,7 @@ public class DefaultAccountAuditLogs implements AccountAuditLogs {
 
     @Override
     public List<AuditLog> getAuditLogsForAccount() {
-        return getAuditLogs(ObjectType.BUNDLE).getAuditLogs(accountId);
+        return getAuditLogs(ObjectType.ACCOUNT).getAuditLogs(accountId);
     }
 
     @Override

diff --git a/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java b/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
index 02d56bb..4fe070e 100644
--- a/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
+++ b/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
@@ -384,9 +384,17 @@ 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 context) {
+    private void insertAudits(final TableName tableName, final Long entityRecordId, final Long historyRecordId, final ChangeType changeType, final InternalCallContext contextMaybeWithoutAccountRecordId) {
         final TableName destinationTableName = Objects.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 = new InternalCallContext(contextMaybeWithoutAccountRecordId, entityRecordId);
+        } else {
+            context = contextMaybeWithoutAccountRecordId;
+        }
         sqlDao.insertAuditFromTransaction(audit, context);
 
         // We need to invalidate the caches. There is a small window of doom here where caches will be stale.