killbill-memoizeit

diff --git a/api/src/main/java/com/ning/billing/util/api/AuditUserApi.java b/api/src/main/java/com/ning/billing/util/api/AuditUserApi.java
index ee9cc2a..340420e 100644
--- a/api/src/main/java/com/ning/billing/util/api/AuditUserApi.java
+++ b/api/src/main/java/com/ning/billing/util/api/AuditUserApi.java
@@ -20,6 +20,7 @@ import java.util.List;
 import java.util.UUID;
 
 import com.ning.billing.ObjectType;
+import com.ning.billing.account.api.AccountApiException;
 import com.ning.billing.entitlement.api.timeline.BundleTimeline;
 import com.ning.billing.entitlement.api.timeline.EntitlementRepairException;
 import com.ning.billing.invoice.api.Invoice;
@@ -27,6 +28,7 @@ import com.ning.billing.invoice.api.InvoicePayment;
 import com.ning.billing.payment.api.Payment;
 import com.ning.billing.payment.api.Refund;
 import com.ning.billing.util.audit.AuditLog;
+import com.ning.billing.util.audit.AuditLogsForAccount;
 import com.ning.billing.util.audit.AuditLogsForBundles;
 import com.ning.billing.util.audit.AuditLogsForInvoicePayments;
 import com.ning.billing.util.audit.AuditLogsForInvoices;
@@ -36,6 +38,8 @@ import com.ning.billing.util.callcontext.TenantContext;
 
 public interface AuditUserApi {
 
+
+    public AuditLogsForAccount getAuditLogsForAccount(UUID accountId, AuditLevel auditLevel, TenantContext context);
     /**
      * Fetch all audit logs for a bundle.
      *

diff --git a/api/src/main/java/com/ning/billing/util/audit/AuditLogsForAccount.java b/api/src/main/java/com/ning/billing/util/audit/AuditLogsForAccount.java
new file mode 100644
index 0000000..005297b
--- /dev/null
+++ b/api/src/main/java/com/ning/billing/util/audit/AuditLogsForAccount.java
@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.util.audit;
+
+import java.util.List;
+
+public interface AuditLogsForAccount {
+
+    List<AuditLog> getAccountAuditLogs();
+}

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/BeatrixIntegrationModule.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/BeatrixIntegrationModule.java
index 79c96bf..27240be 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/BeatrixIntegrationModule.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/BeatrixIntegrationModule.java
@@ -34,7 +34,9 @@ import com.ning.billing.beatrix.extbus.PersistentExternalBus;
 import com.ning.billing.beatrix.integration.overdue.IntegrationTestOverdueModule;
 import com.ning.billing.beatrix.lifecycle.DefaultLifecycle;
 import com.ning.billing.beatrix.lifecycle.Lifecycle;
+import com.ning.billing.beatrix.util.AccountChecker;
 import com.ning.billing.beatrix.util.AuditChecker;
+import com.ning.billing.beatrix.util.EntitlementChecker;
 import com.ning.billing.beatrix.util.InvoiceChecker;
 import com.ning.billing.beatrix.util.PaymentChecker;
 import com.ning.billing.catalog.api.CatalogService;
@@ -116,6 +118,8 @@ public class BeatrixIntegrationModule extends AbstractModule {
         install(new IntegrationTestOverdueModule());
         install(new AuditModule());
 
+        bind(AccountChecker.class).asEagerSingleton();
+        bind(EntitlementChecker.class).asEagerSingleton();
         bind(InvoiceChecker.class).asEagerSingleton();
         bind(PaymentChecker.class).asEagerSingleton();
         bind(AuditChecker.class).asEagerSingleton();

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegration.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegration.java
index d2ae8ea..b1cb705 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegration.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegration.java
@@ -27,6 +27,7 @@ import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
 import com.ning.billing.account.api.Account;
+import com.ning.billing.account.api.AccountData;
 import com.ning.billing.api.TestApiListener.NextEvent;
 import com.ning.billing.beatrix.util.InvoiceChecker.ExpectedInvoiceItemCheck;
 import com.ning.billing.beatrix.util.PaymentChecker.ExpectedPaymentCheck;
@@ -53,28 +54,34 @@ public class TestIntegration extends TestIntegrationBase {
     @Test(groups = "slow")
     public void testCancelBPWithAOTheSameDay() throws Exception {
 
-        final Account account = createAccountWithPaymentMethod(getAccountData(1));
+        final AccountData accountData = getAccountData(1);
+        final Account account = createAccountWithPaymentMethod(accountData);
+        accountChecker.checkAccount(account.getId(), accountData, callContext);
 
         // We take april as it has 30 days (easier to play with BCD)
         // Set clock to the initial start date - we implicitly assume here that the account timezone is UTC
         clock.setDay(new LocalDate(2012, 4, 1));
 
         final SubscriptionBundle bundle = entitlementUserApi.createBundleForAccount(account.getId(), "whatever", callContext);
+        entitlementChecker.checkBundleNoAudits(bundle.getId(), bundle.getAccountId(), bundle.getExternalKey(), callContext);
 
         //
         // CREATE SUBSCRIPTION AND EXPECT BOTH EVENTS: NextEvent.CREATE NextEvent.INVOICE
         //
         final Subscription bpSubscription = createSubscriptionAndCheckForCompletion(bundle.getId(), "Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.INVOICE);
+        // Check bundle after BP got created otherwise we get an error from auditApi.
+        entitlementChecker.checkSubscriptionCreated(bpSubscription.getId(), callContext);
         invoiceChecker.checkInvoice(account.getId(), 1, callContext, new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), null, InvoiceItemType.FIXED, new BigDecimal("0")));
-
+        entitlementChecker.checkBundleAuditUpdated(bundle.getId(), callContext);
         //
         // ADD ADD_ON ON THE SAME DAY
         //
         createSubscriptionAndCheckForCompletion(bundle.getId(), "Telescopic-Scope", ProductCategory.ADD_ON, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.INVOICE, NextEvent.PAYMENT);
         Invoice invoice = invoiceChecker.checkInvoice(account.getId(), 2, callContext, new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.RECURRING, new BigDecimal("399.95")));
         paymentChecker.checkPayment(account.getId(), 1, callContext, new ExpectedPaymentCheck(new LocalDate(2012, 4, 1), new BigDecimal("399.95"), PaymentStatus.SUCCESS, invoice.getId(), Currency.USD));
+        entitlementChecker.checkBundleAuditUpdated(bundle.getId(), callContext);
 
-                                    //
+        //
                                     // CANCEL BP ON THE SAME DAY (we should have two cancellations, BP and AO)
                                     // There is no invoice created as we only adjust the previous invoice.
                                     //
@@ -84,6 +91,8 @@ public class TestIntegration extends TestIntegrationBase {
                                     // The second invoice should be adjusted for the AO (we paid for the full period) and since we paid we should also see a CBA
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 5, 1), InvoiceItemType.REPAIR_ADJ, new BigDecimal("-399.95")),
                                     new ExpectedInvoiceItemCheck(new LocalDate(2012, 4, 1), new LocalDate(2012, 4, 1), InvoiceItemType.CBA_ADJ, new BigDecimal("399.95")));
+        entitlementChecker.checkBundleAuditUpdated(bundle.getId(), callContext);
+
     }
 
     @Test(groups = "slow")

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
index df98f42..8d1962c 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationBase.java
@@ -49,6 +49,8 @@ import com.ning.billing.api.TestListenerStatus;
 import com.ning.billing.beatrix.BeatrixTestSuiteWithEmbeddedDB;
 import com.ning.billing.beatrix.bus.api.ExternalBus;
 import com.ning.billing.beatrix.lifecycle.Lifecycle;
+import com.ning.billing.beatrix.util.AccountChecker;
+import com.ning.billing.beatrix.util.EntitlementChecker;
 import com.ning.billing.beatrix.util.InvoiceChecker;
 import com.ning.billing.beatrix.util.PaymentChecker;
 import com.ning.billing.catalog.api.BillingPeriod;
@@ -183,9 +185,15 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
     protected PaymentChecker paymentChecker;
 
     @Inject
+    protected AccountChecker accountChecker;
+
+    @Inject
     protected ExternalBus externalBus;
 
     @Inject
+    protected EntitlementChecker entitlementChecker;
+
+    @Inject
     protected AccountInternalApi accountInternalApi;
 
     protected TestApiListener busHandler;

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/util/AccountChecker.java b/beatrix/src/test/java/com/ning/billing/beatrix/util/AccountChecker.java
new file mode 100644
index 0000000..e5de4ca
--- /dev/null
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/util/AccountChecker.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.beatrix.util;
+
+import java.util.UUID;
+
+
+import javax.inject.Inject;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+
+import com.ning.billing.account.api.Account;
+import com.ning.billing.account.api.AccountData;
+import com.ning.billing.account.api.AccountUserApi;
+import com.ning.billing.util.callcontext.CallContext;
+
+public class AccountChecker {
+
+    private static final Logger log = LoggerFactory.getLogger(AccountChecker.class);
+
+    private final AccountUserApi accountApi;
+    private final AuditChecker auditChecker;
+
+
+    @Inject
+    public AccountChecker(final AccountUserApi accountApi, final AuditChecker auditChecker) {
+        this.accountApi = accountApi;
+        this.auditChecker = auditChecker;
+    }
+
+
+    public Account checkAccount(final UUID accountId, final AccountData accountData, final CallContext context) throws Exception {
+
+        final Account account = accountApi.getAccountById(accountId, context);
+        // Not all test pass it, since this is always the same test
+        if (accountData != null) {
+            Assert.assertEquals(account.getName(), accountData.getName());
+            Assert.assertEquals(account.getFirstNameLength(), accountData.getFirstNameLength());
+            Assert.assertEquals(account.getEmail(), accountData.getEmail());
+            Assert.assertEquals(account.getPhone(), accountData.getPhone());
+            Assert.assertEquals(account.isNotifiedForInvoices(), accountData.isNotifiedForInvoices());
+            Assert.assertEquals(account.getExternalKey(), accountData.getExternalKey());
+            Assert.assertEquals(account.getBillCycleDay().getDayOfMonthUTC(), accountData.getBillCycleDay().getDayOfMonthUTC());
+            Assert.assertEquals(account.getBillCycleDay().getDayOfMonthLocal(), accountData.getBillCycleDay().getDayOfMonthLocal());
+            Assert.assertEquals(account.getCurrency(), accountData.getCurrency());
+            Assert.assertEquals(account.getTimeZone(), accountData.getTimeZone());
+            // createWithPaymentMethod will update the paymentMethod
+            //Assert.assertEquals(account.getPaymentMethodId(), accountData.getPaymentMethodId());
+        }
+
+        auditChecker.checkAccountCreated(account, context);
+        return account;
+    }
+}

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/util/AuditChecker.java b/beatrix/src/test/java/com/ning/billing/beatrix/util/AuditChecker.java
index 56a17b1..9a7cd01 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/util/AuditChecker.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/util/AuditChecker.java
@@ -30,6 +30,8 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
 
+import com.ning.billing.account.api.Account;
+import com.ning.billing.account.dao.AccountSqlDao;
 import com.ning.billing.entitlement.api.timeline.EntitlementRepairException;
 import com.ning.billing.entitlement.engine.dao.BundleSqlDao;
 import com.ning.billing.entitlement.engine.dao.EntitlementEventSqlDao;
@@ -43,12 +45,14 @@ import com.ning.billing.payment.dao.PaymentSqlDao;
 import com.ning.billing.util.api.AuditLevel;
 import com.ning.billing.util.api.AuditUserApi;
 import com.ning.billing.util.audit.AuditLog;
+import com.ning.billing.util.audit.AuditLogsForAccount;
 import com.ning.billing.util.audit.AuditLogsForBundles;
 import com.ning.billing.util.audit.AuditLogsForInvoices;
 import com.ning.billing.util.audit.AuditLogsForPayments;
 import com.ning.billing.util.audit.ChangeType;
 import com.ning.billing.util.callcontext.CallContext;
 import com.ning.billing.util.callcontext.InternalCallContextFactory;
+import com.ning.billing.util.callcontext.UserType;
 import com.ning.billing.util.entity.Entity;
 import com.ning.billing.util.entity.dao.EntityModelDao;
 import com.ning.billing.util.entity.dao.EntitySqlDao;
@@ -71,61 +75,50 @@ public class AuditChecker {
         this.callContextFactory = callContextFactory;
     }
 
+
+    /***********************************************  ACCOUNT ********************************************************/
+
+    public void checkAccountCreated(final Account account, final CallContext context) {
+        AuditLogsForAccount result = auditUserApi.getAuditLogsForAccount(account.getId(), AuditLevel.FULL, context);
+        checkAuditLog(ChangeType.INSERT, context, result.getAccountAuditLogs().get(0), account.getId(), AccountSqlDao.class, true, true);
+        checkAuditLog(ChangeType.UPDATE, context, result.getAccountAuditLogs().get(1), account.getId(), AccountSqlDao.class, true, true);
+    }
+
+
     public void checkPaymentCreated(final Payment payment, final CallContext context) {
         AuditLogsForPayments result = getAuditLogsForPayment(payment, context);
         Assert.assertEquals(result.getPaymentsAuditLogs().size(), 1);
         final List<AuditLog> paymentLogs = result.getPaymentsAuditLogs().get(payment.getId());
         Assert.assertEquals(paymentLogs.size(), 2);
-        checkAuditLog(ChangeType.INSERT, context, paymentLogs.get(0), payment.getId(), PaymentSqlDao.class, true);
-        checkAuditLog(ChangeType.UPDATE, context, paymentLogs.get(1), payment.getId(), PaymentSqlDao.class, true);
+        checkAuditLog(ChangeType.INSERT, context, paymentLogs.get(0), payment.getId(), PaymentSqlDao.class, true, false);
+        checkAuditLog(ChangeType.UPDATE, context, paymentLogs.get(1), payment.getId(), PaymentSqlDao.class, true, false);
     }
 
-    private AuditLogsForPayments getAuditLogsForPayment(final Payment payment, final CallContext context) {
-        AuditLogsForPayments results = auditUserApi.getAuditLogsForPayments(Collections.singletonList(payment), AuditLevel.FULL, context);
-        return results;
-    }
+    /***********************************************  BUNDLE ********************************************************/
 
-    public void checkInvoiceCreated(final Invoice invoice, final CallContext context) {
-        AuditLogsForInvoices result = getAuditLogForInvoice(invoice, context);
-        Assert.assertEquals(result.getInvoiceAuditLogs().keySet().size(), 1);
-        final List<AuditLog> invoiceLogs = result.getInvoiceAuditLogs().get(invoice.getId());
-        Assert.assertEquals(invoiceLogs.size(), 1);
-        checkAuditLog(ChangeType.INSERT, context, invoiceLogs.get(0), invoice.getId(), InvoiceSqlDao.class, false);
-
-        Assert.assertEquals(result.getInvoiceItemsAuditLogs().keySet().size(), invoice.getInvoiceItems().size());
-        for (InvoiceItem cur : invoice.getInvoiceItems()) {
-            final List<AuditLog> auditLogs = result.getInvoiceItemsAuditLogs().get(cur.getId());
-            Assert.assertEquals(auditLogs.size(), 1);
-            checkAuditLog(ChangeType.INSERT, context, auditLogs.get(0), cur.getId(), InvoiceItemSqlDao.class, false);
-        }
-    }
 
     // Pass the call context used to create the bundle
     public void checkBundleCreated(final UUID bundleId, final CallContext context) {
         final AuditLogsForBundles auditLogsForBundles = getAuditLogsForBundle(bundleId, context);
-
         Assert.assertEquals(auditLogsForBundles.getBundlesAuditLogs().keySet().size(), 1);
-        Assert.assertEquals(auditLogsForBundles.getBundlesAuditLogs().get(bundleId).size(), 1);
-        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getBundlesAuditLogs().get(bundleId).get(0), bundleId, BundleSqlDao.class, false);
+        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getBundlesAuditLogs().get(bundleId).get(0), bundleId, BundleSqlDao.class, false, false);
     }
 
     // Pass the call context used to update the bundle
     public void checkBundleUpdated(final UUID bundleId, final CallContext context) {
         final AuditLogsForBundles auditLogsForBundles = getAuditLogsForBundle(bundleId, context);
-
         Assert.assertEquals(auditLogsForBundles.getBundlesAuditLogs().keySet().size(), 1);
-        Assert.assertEquals(auditLogsForBundles.getBundlesAuditLogs().get(bundleId).size(), 2);
-        checkAuditLog(ChangeType.INSERT, auditLogsForBundles.getBundlesAuditLogs().get(bundleId).get(0));
-        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getBundlesAuditLogs().get(bundleId).get(1), bundleId, BundleSqlDao.class, false);
+        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getBundlesAuditLogs().get(bundleId).get(auditLogsForBundles.getBundlesAuditLogs().get(bundleId).size() - 1), bundleId, BundleSqlDao.class, false, false);
     }
 
+    /***********************************************  SUBSCRIPTION ********************************************************/
+
     // Pass the call context used to create the subscription
     public void checkSubscriptionCreated(final UUID bundleId, final UUID subscriptionId, final CallContext context) {
         final AuditLogsForBundles auditLogsForBundles = getAuditLogsForBundle(bundleId, context);
 
         Assert.assertEquals(auditLogsForBundles.getSubscriptionsAuditLogs().keySet().size(), 1);
-        Assert.assertEquals(auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).size(), 1);
-        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).get(0), subscriptionId, SubscriptionSqlDao.class, false);
+        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).get(0), subscriptionId, SubscriptionSqlDao.class, false, true);
     }
 
     // Pass the call context used to update the subscription
@@ -135,28 +128,53 @@ public class AuditChecker {
         Assert.assertEquals(auditLogsForBundles.getSubscriptionsAuditLogs().keySet().size(), 1);
         Assert.assertEquals(auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).size(), 2);
         checkAuditLog(ChangeType.INSERT, auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).get(0));
-        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).get(1), subscriptionId, SubscriptionSqlDao.class, false);
+        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getSubscriptionsAuditLogs().get(subscriptionId).get(1), subscriptionId, SubscriptionSqlDao.class, false, false);
     }
 
+    /***********************************************  SUBSCRIPTION EVENTS ********************************************************/
+
     // Pass the call context used to create the subscription event
     public void checkSubscriptionEventCreated(final UUID bundleId, final UUID subscriptionEventId, final CallContext context) {
         final AuditLogsForBundles auditLogsForBundles = getAuditLogsForBundle(bundleId, context);
-
-        Assert.assertEquals(auditLogsForBundles.getSubscriptionEventsAuditLogs().keySet().size(), 1);
-        Assert.assertEquals(auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).size(), 1);
-        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).get(0), subscriptionEventId, EntitlementEventSqlDao.class, false);
+        checkAuditLog(ChangeType.INSERT, context, auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).get(0), subscriptionEventId, EntitlementEventSqlDao.class, false, true);
     }
 
     // Pass the call context used to update the subscription event
     public void checkSubscriptionEventUpdated(final UUID bundleId, final UUID subscriptionEventId, final CallContext context) {
         final AuditLogsForBundles auditLogsForBundles = getAuditLogsForBundle(bundleId, context);
-
-        Assert.assertEquals(auditLogsForBundles.getSubscriptionEventsAuditLogs().keySet().size(), 1);
-        Assert.assertEquals(auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).size(), 2);
         checkAuditLog(ChangeType.INSERT, auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).get(0));
-        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).get(1), subscriptionEventId, EntitlementEventSqlDao.class, false);
+        checkAuditLog(ChangeType.UPDATE, context, auditLogsForBundles.getSubscriptionEventsAuditLogs().get(subscriptionEventId).get(1), subscriptionEventId, EntitlementEventSqlDao.class, false, true);
     }
 
+
+
+    /***********************************************  PAYMENT ********************************************************/
+
+    private AuditLogsForPayments getAuditLogsForPayment(final Payment payment, final CallContext context) {
+        AuditLogsForPayments results = auditUserApi.getAuditLogsForPayments(Collections.singletonList(payment), AuditLevel.FULL, context);
+        return results;
+    }
+
+    /***********************************************  INVOICE ********************************************************/
+
+    public void checkInvoiceCreated(final Invoice invoice, final CallContext context) {
+        AuditLogsForInvoices result = getAuditLogForInvoice(invoice, context);
+        Assert.assertEquals(result.getInvoiceAuditLogs().keySet().size(), 1);
+        final List<AuditLog> invoiceLogs = result.getInvoiceAuditLogs().get(invoice.getId());
+        Assert.assertEquals(invoiceLogs.size(), 1);
+        checkAuditLog(ChangeType.INSERT, context, invoiceLogs.get(0), invoice.getId(), InvoiceSqlDao.class, false, false);
+
+        Assert.assertEquals(result.getInvoiceItemsAuditLogs().keySet().size(), invoice.getInvoiceItems().size());
+        for (InvoiceItem cur : invoice.getInvoiceItems()) {
+            final List<AuditLog> auditLogs = result.getInvoiceItemsAuditLogs().get(cur.getId());
+            Assert.assertEquals(auditLogs.size(), 1);
+            checkAuditLog(ChangeType.INSERT, context, auditLogs.get(0), cur.getId(), InvoiceItemSqlDao.class, false, false);
+        }
+    }
+
+
+
+
     private AuditLogsForBundles getAuditLogsForBundle(final UUID bundleId, final CallContext context) {
         try {
             return auditUserApi.getAuditLogsForBundle(bundleId, AuditLevel.FULL, context);
@@ -170,53 +188,27 @@ public class AuditChecker {
         return auditUserApi.getAuditLogsForInvoices(Collections.singletonList(invoice), AuditLevel.FULL, context);
     }
 
+
     private void checkAuditLog(final ChangeType insert, final AuditLog auditLog) {
-        checkAuditLog(insert, null, auditLog, null, EntitySqlDao.class, false);
+        checkAuditLog(insert, null, auditLog, null, EntitySqlDao.class, false, false);
     }
 
 
-    private <T extends EntitySqlDao<M, E>, M extends EntityModelDao<E>, E extends Entity> void checkAuditLog(final ChangeType changeType, @Nullable final CallContext context, final AuditLog auditLog, final UUID entityId, Class<T> sqlDao, boolean useHistory) {
+    private <T extends EntitySqlDao<M, E>, M extends EntityModelDao<E>, E extends Entity> void checkAuditLog(final ChangeType changeType, @Nullable final CallContext context, final AuditLog auditLog, final UUID entityId, Class<T> sqlDao,
+                                                                                                             boolean useHistory, boolean checkContext) {
         Assert.assertEquals(auditLog.getChangeType(), changeType);
 
-        if (context != null) {
-            /*
+        if (checkContext) {
             Assert.assertEquals(auditLog.getUserName(), context.getUserName());
-            Assert.assertEquals(auditLog.getReasonCode(), context.getReasonCode());
-            // TODO check 'Next Billing Date' and 'Transition' - add comment, maybe internal reason code and token
             Assert.assertEquals(auditLog.getComment(), context.getComments());
-            Assert.assertEquals(auditLog.getUserToken(), context.getUserToken().toString());
-            Assert.assertEquals(auditLog.getCreatedDate(), context.getCreatedDate());
-            */
-
-            final M entityModel = extractEntityModelFromEntityWithTargetRecordId(auditLog.getId(), sqlDao, context, useHistory);
-            Assert.assertEquals(entityModel.getId(), entityId);
+            //Assert.assertEquals(auditLog.getCreatedDate().comparesTo(context.getCreatedDate()));
         }
-    }
-
-    /*
-    private <T extends EntitySqlDao<M, E>, M extends EntityModelDao<E>, E extends Entity> UUID extractEntityIdFromModel(M entityModel, CallContext context, Class<T> sqlDao) {
-
-        final UUID entityModelId = entityModel.getId();
-
-        if ((entityModel.getTableName().equals(TableName.PAYMENTS)) ||
-             (entityModel.getTableName().equals(TableName.PAYMENT_ATTEMPTS)) ||
-             (entityModel.getTableName().equals(TableName.PAYMENT_METHODS)) ||
-             (entityModel.getTableName().equals(TableName.REFUNDS)) ||
-             (entityModel.getTableName().equals(TableName.ACCOUNT)) ||
-             (entityModel.getTableName().equals(TableName.ACCOUNT_EMAIL)) ||
-             (entityModel.getTableName().equals(TableName.TAG)) ||
-             (entityModel.getTableName().equals(TableName.TAG_DEFINITIONS)) ||
-             (entityModel.getTableName().equals(TableName.CUSTOM_FIELD))) {
+        Assert.assertEquals(auditLog.getUserToken(), context.getUserToken().toString());
+        final M entityModel = extractEntityModelFromEntityWithTargetRecordId(auditLog.getId(), sqlDao, context, useHistory);
+        Assert.assertEquals(entityModel.getId(), entityId);
 
-            M historyEntityModel = extractEntityModelFromEntityWithTargetRecordId(entityModelId, sqlDao, context);
-
-            return extractEntityIdFromModel(historyEntityModel, context, sqlDao);
-        } else {
-            return entityModelId;
-        }
     }
 
-    */
 
     private <T extends EntitySqlDao<M, E>, M extends EntityModelDao<E>, E extends Entity> M extractEntityModelFromEntityWithTargetRecordId(final UUID entityId, final Class<T> sqlDao, final CallContext context, final boolean useHistory) {
         Integer targetRecordId = dbi.withHandle(new HandleCallback<Integer>() {
@@ -234,4 +226,5 @@ public class AuditChecker {
         }
         return dbi.onDemand(sqlDao).getByRecordId(Long.valueOf(targetRecordId), callContextFactory.createInternalCallContext(context));
     }
+
 }

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/util/EntitlementChecker.java b/beatrix/src/test/java/com/ning/billing/beatrix/util/EntitlementChecker.java
new file mode 100644
index 0000000..b2c0d02
--- /dev/null
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/util/EntitlementChecker.java
@@ -0,0 +1,84 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.beatrix.util;
+
+import java.util.List;
+import java.util.UUID;
+
+import javax.inject.Inject;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+
+import com.ning.billing.account.api.AccountUserApi;
+import com.ning.billing.entitlement.api.user.EntitlementUserApi;
+import com.ning.billing.entitlement.api.user.EntitlementUserApiException;
+import com.ning.billing.entitlement.api.user.Subscription;
+import com.ning.billing.entitlement.api.user.SubscriptionBundle;
+import com.ning.billing.entitlement.api.user.SubscriptionData;
+import com.ning.billing.entitlement.api.user.SubscriptionTransitionData;
+import com.ning.billing.junction.plumbing.api.BlockingSubscription;
+import com.ning.billing.util.callcontext.CallContext;
+
+public class EntitlementChecker {
+
+
+    private static final Logger log = LoggerFactory.getLogger(EntitlementChecker.class);
+
+    private final EntitlementUserApi entitlementApi;
+    private final AuditChecker auditChecker;
+
+    @Inject
+    public EntitlementChecker(final EntitlementUserApi entitlementApi, final AuditChecker auditChecker) {
+        this.entitlementApi = entitlementApi;
+        this.auditChecker = auditChecker;
+    }
+
+    public SubscriptionBundle checkBundleNoAudits(final UUID bundleId, final UUID expectedAccountId, final String expectedKey, final CallContext context) throws EntitlementUserApiException {
+        final SubscriptionBundle bundle =  entitlementApi.getBundleFromId(bundleId, context);
+        Assert.assertNotNull(bundle);
+        Assert.assertEquals(bundle.getAccountId(), expectedAccountId);
+        Assert.assertEquals(bundle.getExternalKey(), expectedKey);
+        return bundle;
+    }
+
+    public SubscriptionBundle checkBundleAuditUpdated(final UUID bundleId, final CallContext context) throws EntitlementUserApiException {
+        final SubscriptionBundle bundle =  entitlementApi.getBundleFromId(bundleId, context);
+        auditChecker.checkBundleUpdated(bundle.getId(), context);
+        return bundle;
+    }
+
+    public Subscription checkSubscriptionCreated(final UUID subscriptionId, final CallContext context) throws EntitlementUserApiException {
+        final Subscription subscription = entitlementApi.getSubscriptionFromId(subscriptionId, context);
+        Assert.assertNotNull(subscription);
+        auditChecker.checkSubscriptionCreated(subscription.getBundleId(), subscriptionId, context);
+
+        List<SubscriptionTransitionData> subscriptionEvents = getSubscriptionEvents(subscription);
+        Assert.assertTrue(subscriptionEvents.size() >= 1);
+        auditChecker.checkSubscriptionEventCreated(subscription.getBundleId(), subscriptionEvents.get(0).getId(), context);
+
+        auditChecker.checkBundleCreated(subscription.getBundleId(), context);
+        return subscription;
+    }
+
+    private List<SubscriptionTransitionData> getSubscriptionEvents(final Subscription subscription) {
+        return ((SubscriptionData) ((BlockingSubscription) subscription).getDelegateSubscription()).getAllTransitions();
+    }
+
+
+}

diff --git a/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditLogsForAccount.java b/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditLogsForAccount.java
new file mode 100644
index 0000000..388e668
--- /dev/null
+++ b/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditLogsForAccount.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2012 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.ning.billing.util.audit.api;
+
+import java.util.List;
+
+import com.ning.billing.util.audit.AuditLog;
+import com.ning.billing.util.audit.AuditLogsForAccount;
+
+public class DefaultAuditLogsForAccount implements AuditLogsForAccount {
+
+    private final List<AuditLog> auditLogsForAccount;
+
+    public DefaultAuditLogsForAccount(final List<AuditLog> auditLogsForAccount) {
+        this.auditLogsForAccount = auditLogsForAccount;
+    }
+
+    @Override
+    public List<AuditLog> getAccountAuditLogs() {
+        return auditLogsForAccount;
+    }
+}

diff --git a/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditUserApi.java b/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditUserApi.java
index e95e1b2..3cb03c2 100644
--- a/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditUserApi.java
+++ b/util/src/main/java/com/ning/billing/util/audit/api/DefaultAuditUserApi.java
@@ -24,6 +24,7 @@ import java.util.UUID;
 import javax.inject.Inject;
 
 import com.ning.billing.ObjectType;
+import com.ning.billing.account.api.AccountApiException;
 import com.ning.billing.entitlement.api.timeline.BundleTimeline;
 import com.ning.billing.entitlement.api.timeline.EntitlementRepairException;
 import com.ning.billing.entitlement.api.timeline.EntitlementTimelineApi;
@@ -37,6 +38,7 @@ import com.ning.billing.payment.api.Refund;
 import com.ning.billing.util.api.AuditLevel;
 import com.ning.billing.util.api.AuditUserApi;
 import com.ning.billing.util.audit.AuditLog;
+import com.ning.billing.util.audit.AuditLogsForAccount;
 import com.ning.billing.util.audit.AuditLogsForBundles;
 import com.ning.billing.util.audit.AuditLogsForInvoicePayments;
 import com.ning.billing.util.audit.AuditLogsForInvoices;
@@ -68,6 +70,11 @@ public class DefaultAuditUserApi implements AuditUserApi {
     }
 
     @Override
+    public AuditLogsForAccount getAuditLogsForAccount(final UUID accountId, final AuditLevel auditLevel, final TenantContext context){
+        return new DefaultAuditLogsForAccount(getAuditLogs(accountId, ObjectType.ACCOUNT, auditLevel, context));
+    }
+
+    @Override
     public AuditLogsForBundles getAuditLogsForBundle(final UUID bundleId, final AuditLevel auditLevel, final TenantContext context) throws EntitlementRepairException {
         return getAuditLogsForBundles(ImmutableList.<BundleTimeline>of(timelineApi.getBundleTimeline(bundleId, context)), auditLevel, context);
     }

diff --git a/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoStringTemplate.java b/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoStringTemplate.java
index 49e85c9..f96ea42 100644
--- a/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoStringTemplate.java
+++ b/util/src/main/java/com/ning/billing/util/entity/dao/EntitySqlDaoStringTemplate.java
@@ -75,7 +75,7 @@ public @interface EntitySqlDaoStringTemplate {
                                 if (sqlObjectType.getGenericInterfaces()[i] instanceof ParameterizedType) {
                                     final ParameterizedType type = (ParameterizedType) sqlObjectType.getGenericInterfaces()[i];
                                     for (int j = 0; j < type.getActualTypeArguments().length; j++) {
-                                        final Class modelClazz = (Class) type.getActualTypeArguments()[i];
+                                        final Class modelClazz = (Class) type.getActualTypeArguments()[j];
                                         if (Entity.class.isAssignableFrom(modelClazz)) {
                                             query.registerMapper(new LowerToCamelBeanMapperFactory(modelClazz));
                                         }

diff --git a/util/src/main/resources/com/ning/billing/util/entity/dao/EntitySqlDao.sql.stg b/util/src/main/resources/com/ning/billing/util/entity/dao/EntitySqlDao.sql.stg
index 9a0738f..6610aa5 100644
--- a/util/src/main/resources/com/ning/billing/util/entity/dao/EntitySqlDao.sql.stg
+++ b/util/src/main/resources/com/ning/billing/util/entity/dao/EntitySqlDao.sql.stg
@@ -292,7 +292,7 @@ from <auditTableName()> t
 where t.target_record_id = :targetRecordId
 and t.table_name = :tableName
 <AND_CHECK_TENANT("t.")>
-order by created_date ASC
+order by <recordIdField("t.")> ASC
 ;
 >>