killbill-aplcache

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 2456676..78f7633 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
@@ -32,6 +32,9 @@ import org.testng.Assert;
 
 import com.ning.billing.account.api.Account;
 import com.ning.billing.account.dao.AccountSqlDao;
+import com.ning.billing.entitlement.api.SubscriptionApi;
+import com.ning.billing.entitlement.api.SubscriptionApiException;
+import com.ning.billing.entitlement.api.SubscriptionBundle;
 import com.ning.billing.subscription.engine.dao.BundleSqlDao;
 import com.ning.billing.subscription.engine.dao.SubscriptionEventSqlDao;
 import com.ning.billing.subscription.engine.dao.SubscriptionSqlDao;
@@ -56,6 +59,7 @@ import com.ning.billing.util.entity.Entity;
 import com.ning.billing.util.entity.dao.EntityModelDao;
 import com.ning.billing.util.entity.dao.EntitySqlDao;
 
+import com.google.common.collect.Lists;
 import com.google.inject.Inject;
 
 public class AuditChecker {
@@ -64,14 +68,16 @@ public class AuditChecker {
 
     private final AuditUserApi auditUserApi;
     private final IDBI dbi;
+    private final SubscriptionApi subscriptionApi;
     private final InternalCallContextFactory callContextFactory;
     private final NonEntityDao nonEntityDao;
 
 
     @Inject
-    public AuditChecker(final AuditUserApi auditUserApi, final IDBI dbi, final InternalCallContextFactory callContextFactory, final NonEntityDao nonEntityDao) {
+    public AuditChecker(final AuditUserApi auditUserApi, final IDBI dbi, final SubscriptionApi subscriptionApi, final InternalCallContextFactory callContextFactory, final NonEntityDao nonEntityDao) {
         this.auditUserApi = auditUserApi;
         this.dbi = dbi;
+        this.subscriptionApi = subscriptionApi;
         this.callContextFactory = callContextFactory;
         this.nonEntityDao = nonEntityDao;
     }
@@ -186,12 +192,13 @@ public class AuditChecker {
 
 
     private AuditLogsForBundles getAuditLogsForBundle(final UUID bundleId, final CallContext context) {
-       // try {
-            return auditUserApi.getAuditLogsForBundle(bundleId, AuditLevel.FULL, context);
-        //} catch (SubscriptionBaseRepairException e) {
-        //    Assert.fail(e.toString());
-        //    return null;
-       // }
+        try {
+            final SubscriptionBundle bundle = subscriptionApi.getSubscriptionBundle(bundleId, context);
+            return auditUserApi.getAuditLogsForBundles(Collections.singletonList(bundle), AuditLevel.FULL, context);
+        } catch (SubscriptionApiException e) {
+            Assert.fail(e.toString());
+            return null;
+        }
     }
 
     private AuditLogsForInvoices getAuditLogForInvoice(final Invoice invoice, final CallContext context) {

diff --git a/jaxrs/src/main/java/com/ning/billing/jaxrs/json/AccountTimelineJson.java b/jaxrs/src/main/java/com/ning/billing/jaxrs/json/AccountTimelineJson.java
index 09ddbed..3c52637 100644
--- a/jaxrs/src/main/java/com/ning/billing/jaxrs/json/AccountTimelineJson.java
+++ b/jaxrs/src/main/java/com/ning/billing/jaxrs/json/AccountTimelineJson.java
@@ -27,7 +27,6 @@ import java.util.UUID;
 import javax.annotation.Nullable;
 
 import com.ning.billing.account.api.Account;
-import com.ning.billing.entitlement.api.Subscription;
 import com.ning.billing.entitlement.api.SubscriptionBundle;
 import com.ning.billing.invoice.api.Invoice;
 import com.ning.billing.invoice.api.InvoiceItem;
@@ -35,7 +34,6 @@ import com.ning.billing.invoice.api.InvoiceItemType;
 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.subscription.api.timeline.BundleBaseTimeline;
 import com.ning.billing.util.audit.AuditLog;
 import com.ning.billing.util.audit.AuditLogsForBundles;
 import com.ning.billing.util.audit.AuditLogsForInvoicePayments;
@@ -98,7 +96,6 @@ public class AccountTimelineJson {
     }
 
 
-
     public AccountTimelineJson(final Account account, final List<Invoice> invoices, final List<Payment> payments,
                                final List<SubscriptionBundle> bundlesTimeline, final Multimap<UUID, Refund> refundsByPayment,
                                final Multimap<UUID, InvoicePayment> chargebacksByPayment, @Nullable final AuditLogsForInvoices invoicesAuditLogs,

diff --git a/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java b/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
index 358a754..1dc8dcd 100644
--- a/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
+++ b/jaxrs/src/main/java/com/ning/billing/jaxrs/resources/AccountResource.java
@@ -19,6 +19,7 @@ package com.ning.billing.jaxrs.resources;
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.UUID;
 
@@ -72,14 +73,17 @@ import com.ning.billing.payment.api.PaymentApi;
 import com.ning.billing.payment.api.PaymentApiException;
 import com.ning.billing.payment.api.PaymentMethod;
 import com.ning.billing.payment.api.Refund;
+import com.ning.billing.subscription.api.timeline.BundleBaseTimeline;
 import com.ning.billing.subscription.api.timeline.SubscriptionBaseRepairException;
 import com.ning.billing.subscription.api.timeline.SubscriptionBaseTimelineApi;
+import com.ning.billing.subscription.api.user.SubscriptionBaseBundle;
 import com.ning.billing.util.api.AuditUserApi;
 import com.ning.billing.util.api.CustomFieldApiException;
 import com.ning.billing.util.api.CustomFieldUserApi;
 import com.ning.billing.util.api.TagApiException;
 import com.ning.billing.util.api.TagDefinitionApiException;
 import com.ning.billing.util.api.TagUserApi;
+import com.ning.billing.util.audit.AuditLogsForBundles;
 import com.ning.billing.util.audit.AuditLogsForInvoicePayments;
 import com.ning.billing.util.audit.AuditLogsForInvoices;
 import com.ning.billing.util.audit.AuditLogsForPayments;
@@ -250,7 +254,7 @@ public class AccountResource extends JaxRsResourceBase {
     @Produces(APPLICATION_JSON)
     public Response getAccountTimeline(@PathParam("accountId") final String accountIdString,
                                        @QueryParam(QUERY_AUDIT) @DefaultValue("NONE") final AuditMode auditMode,
-                                       @javax.ws.rs.core.Context final HttpServletRequest request) throws AccountApiException, PaymentApiException, SubscriptionBaseRepairException {
+                                       @javax.ws.rs.core.Context final HttpServletRequest request) throws AccountApiException, PaymentApiException, SubscriptionApiException {
         final TenantContext tenantContext = context.createContext(request);
 
         final UUID accountId = UUID.fromString(accountIdString);
@@ -280,22 +284,14 @@ public class AccountResource extends JaxRsResourceBase {
             chargebacksByPayment.put(chargeback.getPaymentId(), chargeback);
         }
 
-        /*
         // Get the bundles
-        final List<SubscriptionBaseBundle> bundles = entitlementApi.getBundlesForAccount(account.getId(), tenantContext);
-        final List<BundleBaseTimeline> bundlesTimeline = new LinkedList<BundleBaseTimeline>();
-        for (final SubscriptionBaseBundle bundle : bundles) {
-            bundlesTimeline.add(timelineApi.getBundleTimeline(bundle.getId(), tenantContext));
-        }
-        final AuditLogsForBundles bundlesAuditLogs = auditUserApi.getAuditLogsForBundles(bundlesTimeline, auditMode.getLevel(), tenantContext);
+        final List<SubscriptionBundle> bundles = subscriptionApi.getSubscriptionBundlesForAccountId(account.getId(), tenantContext);
+        final AuditLogsForBundles bundlesAuditLogs = auditUserApi.getAuditLogsForBundles(bundles, auditMode.getLevel(), tenantContext);
 
-        final AccountTimelineJson json = new AccountTimelineJson(account, invoices, payments, bundlesTimeline,
+        final AccountTimelineJson json = new AccountTimelineJson(account, invoices, payments, bundles,
                                                                  refundsByPayment, chargebacksByPayment,
                                                                  invoicesAuditLogs, paymentsAuditLogs, refundsAuditLogs,
                                                                  chargebacksAuditLogs, bundlesAuditLogs);
-        // STEPH_ENT
-        */
-        final AccountTimelineJson json = null;
         return Response.status(Status.OK).entity(json).build();
     }
 

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 6600e92..5925428 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
@@ -21,9 +21,13 @@ import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import com.ning.billing.ObjectType;
+import com.ning.billing.entitlement.api.Subscription;
+import com.ning.billing.entitlement.api.SubscriptionBundle;
+import com.ning.billing.entitlement.api.SubscriptionBundleTimeline.SubscriptionEvent;
 import com.ning.billing.subscription.api.timeline.BundleBaseTimeline;
 import com.ning.billing.subscription.api.timeline.SubscriptionBaseRepairException;
 import com.ning.billing.subscription.api.timeline.SubscriptionBaseTimeline;
@@ -53,18 +57,18 @@ import com.ning.billing.util.callcontext.InternalCallContextFactory;
 import com.ning.billing.util.callcontext.TenantContext;
 import com.ning.billing.util.dao.TableName;
 
+import com.google.common.base.Predicate;
+import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
 
 public class DefaultAuditUserApi implements AuditUserApi {
 
     private final AuditDao auditDao;
-    private final SubscriptionBaseTimelineApi timelineApi;
     private final InternalCallContextFactory internalCallContextFactory;
 
     @Inject
     public DefaultAuditUserApi(final AuditDao auditDao, final SubscriptionBaseTimelineApi timelineApi, final InternalCallContextFactory internalCallContextFactory) {
         this.auditDao = auditDao;
-        this.timelineApi = timelineApi;
         this.internalCallContextFactory = internalCallContextFactory;
     }
 
@@ -74,27 +78,23 @@ public class DefaultAuditUserApi implements AuditUserApi {
     }
 
     @Override
-    public AuditLogsForBundles getAuditLogsForBundle(final UUID bundleId, final AuditLevel auditLevel, final TenantContext context)  {
-
-        try {
-            return getAuditLogsForBundles(ImmutableList.<BundleBaseTimeline>of(timelineApi.getBundleTimeline(bundleId, context)), auditLevel, context);
-        } catch (SubscriptionBaseRepairException e) {
-            // STEPH_ENT
-            e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
-            return null;
-        }
-    }
-
-    public AuditLogsForBundles getAuditLogsForBundles(final List<BundleBaseTimeline> bundles, final AuditLevel auditLevel, final TenantContext context) {
+    public AuditLogsForBundles getAuditLogsForBundles(final List<SubscriptionBundle> bundles, final AuditLevel auditLevel, final TenantContext context) {
         final Map<UUID, List<AuditLog>> bundlesAuditLogs = new HashMap<UUID, List<AuditLog>>();
         final Map<UUID, List<AuditLog>> subscriptionsAuditLogs = new HashMap<UUID, List<AuditLog>>();
         final Map<UUID, List<AuditLog>> subscriptionEventsAuditLogs = new HashMap<UUID, List<AuditLog>>();
-        for (final BundleBaseTimeline bundle : bundles) {
+        for (final SubscriptionBundle bundle : bundles) {
             bundlesAuditLogs.put(bundle.getId(), getAuditLogs(bundle.getId(), ObjectType.BUNDLE, auditLevel, context));
-            for (final SubscriptionBaseTimeline subscriptionTimeline : bundle.getSubscriptions()) {
-                subscriptionsAuditLogs.put(subscriptionTimeline.getId(), getAuditLogs(subscriptionTimeline.getId(), ObjectType.SUBSCRIPTION, auditLevel, context));
-                for (final ExistingEvent event : subscriptionTimeline.getExistingEvents()) {
-                    subscriptionEventsAuditLogs.put(event.getEventId(), getAuditLogs(event.getEventId(), ObjectType.SUBSCRIPTION_EVENT, auditLevel, context));
+            for (final Subscription cur : bundle.getSubscriptions()) {
+
+                final ImmutableList<SubscriptionEvent> events =  ImmutableList.<SubscriptionEvent>copyOf(Collections2.filter(bundle.getTimeline().getSubscriptionEvents(), new Predicate<SubscriptionEvent>() {
+                    @Override
+                    public boolean apply(@Nullable final SubscriptionEvent input) {
+                        return input.getEntitlementId().equals(cur.getId());
+                    }
+                }));
+                subscriptionsAuditLogs.put(cur.getId(), getAuditLogs(cur.getId(), ObjectType.SUBSCRIPTION, auditLevel, context));
+                for (final SubscriptionEvent event : events) {
+                    subscriptionEventsAuditLogs.put(event.getId(), getAuditLogs(event.getId(), event.getSubscriptionEventType().getObjectType(), auditLevel, context));
                 }
             }
         }