killbill-uncached

Fixes #310

4/13/2015 6:16:48 PM

Details

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/api/migration/DefaultInvoiceMigrationApi.java b/invoice/src/main/java/org/killbill/billing/invoice/api/migration/DefaultInvoiceMigrationApi.java
index 57ef684..78696a0 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/api/migration/DefaultInvoiceMigrationApi.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/api/migration/DefaultInvoiceMigrationApi.java
@@ -21,7 +21,11 @@ import java.util.List;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
+import org.killbill.billing.account.api.Account;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
+import org.killbill.billing.util.timezone.DateAndTimeZoneContext;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -64,8 +68,9 @@ public class DefaultInvoiceMigrationApi implements InvoiceMigrationApi {
 
     @Override
     public UUID createMigrationInvoice(final UUID accountId, final LocalDate targetDate, final BigDecimal balance, final Currency currency, final CallContext context) {
+        Account account;
         try {
-            accountUserApi.getAccountById(accountId, internalCallContextFactory.createInternalTenantContext(accountId, context));
+            account = accountUserApi.getAccountById(accountId, internalCallContextFactory.createInternalTenantContext(accountId, context));
         } catch (AccountApiException e) {
             log.warn("Unable to find account for id {}", accountId);
             return null;
@@ -75,8 +80,11 @@ public class DefaultInvoiceMigrationApi implements InvoiceMigrationApi {
         final InvoiceItemModelDao migrationInvoiceItem = new InvoiceItemModelDao(context.getCreatedDate(), InvoiceItemType.FIXED, migrationInvoice.getId(), accountId, null, null,
                                                                                  null, MigrationPlan.MIGRATION_PLAN_NAME, MigrationPlan.MIGRATION_PLAN_PHASE_NAME, null,
                                                                                  targetDate, null, balance, null, currency, null);
+
+        final DateTime wrongEffectiveDateButDoesNotMatter = null;
+        final DateAndTimeZoneContext dateAndTimeZoneContext = new DateAndTimeZoneContext(wrongEffectiveDateButDoesNotMatter, account.getTimeZone(), clock);
         dao.createInvoice(migrationInvoice, ImmutableList.<InvoiceItemModelDao>of(migrationInvoiceItem),
-                          true, ImmutableMap.<UUID, List<DateTime>>of(), internalCallContextFactory.createInternalCallContext(accountId, context));
+                          true, new FutureAccountNotifications(dateAndTimeZoneContext, ImmutableMap.<UUID, List<DateTime>>of()), internalCallContextFactory.createInternalCallContext(accountId, context));
 
         return migrationInvoice.getId();
     }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
index 762f586..022da79 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
@@ -35,6 +35,7 @@ import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.entity.EntityPersistenceException;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
 import org.killbill.billing.invoice.api.InvoiceItemType;
@@ -217,7 +218,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
 
     @Override
     public void createInvoice(final InvoiceModelDao invoice, final List<InvoiceItemModelDao> invoiceItems,
-                              final boolean isRealInvoice, final Map<UUID, List<DateTime>> callbackDateTimePerSubscriptions,
+                              final boolean isRealInvoice, final FutureAccountNotifications callbackDateTimePerSubscriptions,
                               final InternalCallContext context) {
         transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<Void>() {
             @Override
@@ -775,19 +776,19 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
     }
 
     private void notifyOfFutureBillingEvents(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                             final Map<UUID, List<DateTime>> callbackDateTimePerSubscriptions, final InternalCallContext internalCallContext) {
+                                             final FutureAccountNotifications callbackDateTimePerSubscriptions, final InternalCallContext internalCallContext) {
 
         final long dryRunNotificationTime = invoiceConfig.getDryRunNotificationSchedule().getMillis();
         final boolean isInvoiceNotificationEnabled = dryRunNotificationTime > 0;
-        for (final UUID subscriptionId : callbackDateTimePerSubscriptions.keySet()) {
-            final List<DateTime> callbackDateTimeUTC = callbackDateTimePerSubscriptions.get(subscriptionId);
+        for (final UUID subscriptionId : callbackDateTimePerSubscriptions.getNotifications().keySet()) {
+            final List<DateTime> callbackDateTimeUTC = callbackDateTimePerSubscriptions.getNotifications().get(subscriptionId);
             for (final DateTime cur : callbackDateTimeUTC) {
                 if (isInvoiceNotificationEnabled) {
                     final DateTime curDryRunNotificationTime = cur.minus(dryRunNotificationTime);
                     final DateTime effectiveCurDryRunNotificationTime = (curDryRunNotificationTime.isAfter(clock.getUTCNow())) ? curDryRunNotificationTime : clock.getUTCNow();
-                    nextBillingDatePoster.insertNextBillingDryRunNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, effectiveCurDryRunNotificationTime, cur, internalCallContext);
+                    nextBillingDatePoster.insertNextBillingDryRunNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, effectiveCurDryRunNotificationTime, cur, callbackDateTimePerSubscriptions.getAccountDateAndTimeZoneContext(), internalCallContext);
                 }
-                nextBillingDatePoster.insertNextBillingNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, cur, internalCallContext);
+                nextBillingDatePoster.insertNextBillingNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, cur, callbackDateTimePerSubscriptions.getAccountDateAndTimeZoneContext(), internalCallContext);
             }
         }
     }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDao.java
index 7f2ed99..f8612fa 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/InvoiceDao.java
@@ -30,6 +30,7 @@ import org.joda.time.LocalDate;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
 import org.killbill.billing.util.entity.Pagination;
@@ -38,7 +39,7 @@ import org.killbill.billing.util.entity.dao.EntityDao;
 public interface InvoiceDao extends EntityDao<InvoiceModelDao, Invoice, InvoiceApiException> {
 
     void createInvoice(final InvoiceModelDao invoice, final List<InvoiceItemModelDao> invoiceItems,
-                       final boolean isRealInvoice, final Map<UUID, List<DateTime>> callbackDateTimePerSubscriptions,
+                       final boolean isRealInvoice, final FutureAccountNotifications callbackDateTimePerSubscriptions,
                        final InternalCallContext context);
 
     List<InvoiceItemModelDao> createInvoices(final List<InvoiceModelDao> invoices, final InternalCallContext context);
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java b/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
index 9116ba0..70a28e7 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
@@ -226,7 +226,7 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
         }
 
         // Pretty-print the generated invoice items from the junction events
-        final StringBuilder logStringBuilder = new StringBuilder("Invoice items generated for invoiceId ")
+        final StringBuilder logStringBuilder = new StringBuilder("Proposed Invoice items for invoiceId ")
                 .append(invoiceId)
                 .append(" and accountId ")
                 .append(accountId);
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
index 11a0813..f7c4c4f 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
@@ -31,6 +31,7 @@ import java.util.UUID;
 import javax.annotation.Nullable;
 
 import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
 import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.Account;
@@ -287,8 +288,8 @@ public class InvoiceDispatcher {
                                                                                                                        }
                                                                                                                    }));
 
-                final Map<UUID, List<DateTime>> callbackDateTimePerSubscriptions = createNextFutureNotificationDate(invoiceItemModelDaos, billingEvents.getUsages(), dateAndTimeZoneContext);
-                invoiceDao.createInvoice(invoiceModelDao, invoiceItemModelDaos, isRealInvoiceWithItems, callbackDateTimePerSubscriptions, context);
+                final FutureAccountNotifications futureAccountNotifications = createNextFutureNotificationDate(invoiceItemModelDaos, billingEvents.getUsages(), dateAndTimeZoneContext);
+                invoiceDao.createInvoice(invoiceModelDao, invoiceItemModelDaos, isRealInvoiceWithItems, futureAccountNotifications, context);
 
                 final List<InvoiceItem> fixedPriceInvoiceItems = invoice.getInvoiceItems(FixedPriceInvoiceItem.class);
                 final List<InvoiceItem> recurringInvoiceItems = invoice.getInvoiceItems(RecurringInvoiceItem.class);
@@ -353,8 +354,9 @@ public class InvoiceDispatcher {
         return internalCallContextFactory.createCallContext(context);
     }
 
+
     @VisibleForTesting
-    Map<UUID, List<DateTime>> createNextFutureNotificationDate(final List<InvoiceItemModelDao> invoiceItems, final Map<String, Usage> knownUsages, final DateAndTimeZoneContext dateAndTimeZoneContext) {
+    FutureAccountNotifications createNextFutureNotificationDate(final List<InvoiceItemModelDao> invoiceItems, final Map<String, Usage> knownUsages, final DateAndTimeZoneContext dateAndTimeZoneContext) {
 
         final Map<UUID, List<DateTime>> result = new HashMap<UUID, List<DateTime>>();
 
@@ -405,7 +407,7 @@ public class InvoiceDispatcher {
             perSubscriptionCallback.add(subscriptionUsageCallbackDate);
         }
 
-        return result;
+        return new FutureAccountNotifications(dateAndTimeZoneContext, result);
     }
 
     private DateTime getNextUsageBillingDate(final String usageName, final LocalDate chargedThroughDate, final DateAndTimeZoneContext dateAndTimeZoneContext, final Map<String, Usage> knownUsages) {
@@ -457,6 +459,25 @@ public class InvoiceDispatcher {
         }
     }
 
+    public static class FutureAccountNotifications {
+        private final DateAndTimeZoneContext accountDateAndTimeZoneContext;
+        private final Map<UUID, List<DateTime>> notifications;
+
+        public FutureAccountNotifications(final DateAndTimeZoneContext accountDateAndTimeZoneContext, final Map<UUID, List<DateTime>> notifications) {
+            this.accountDateAndTimeZoneContext = accountDateAndTimeZoneContext;
+            this.notifications = notifications;
+        }
+
+        public DateAndTimeZoneContext getAccountDateAndTimeZoneContext() {
+            return accountDateAndTimeZoneContext;
+        }
+
+        public Map<UUID, List<DateTime>> getNotifications() {
+            return notifications;
+        }
+    }
+
+
     private final static class NullDryRunArguments implements DryRunArguments {
         @Override
         public PlanPhaseSpecifier getPlanPhaseSpecifier() {
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java b/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
index e76a1f5..4a04504 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
@@ -19,18 +19,24 @@
 package org.killbill.billing.invoice.notification;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
+import org.joda.time.LocalDate;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.invoice.api.DefaultInvoiceService;
 import org.killbill.billing.util.entity.dao.EntitySqlDaoWrapperFactory;
+import org.killbill.billing.util.timezone.DateAndTimeZoneContext;
+import org.killbill.notificationq.api.NotificationEventWithMetadata;
 import org.killbill.notificationq.api.NotificationQueue;
 import org.killbill.notificationq.api.NotificationQueueService;
 import org.killbill.notificationq.api.NotificationQueueService.NoSuchNotificationQueue;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Predicate;
+import com.google.common.collect.Iterables;
 import com.google.inject.Inject;
 
 public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
@@ -46,29 +52,52 @@ public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
 
     @Override
     public void insertNextBillingNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                             final UUID subscriptionId, final DateTime futureNotificationTime, final InternalCallContext internalCallContext) {
-        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, accountId, subscriptionId, Boolean.FALSE, futureNotificationTime, futureNotificationTime, internalCallContext);
+                                                             final UUID subscriptionId, final DateTime futureNotificationTime, final DateAndTimeZoneContext accountDateAndTimeZoneContext, final InternalCallContext internalCallContext) {
+        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, accountId, subscriptionId, Boolean.FALSE, futureNotificationTime, futureNotificationTime, accountDateAndTimeZoneContext, internalCallContext);
     }
 
 
     @Override
     public void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                                   final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext) {
-        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, accountId, subscriptionId, Boolean.TRUE, futureNotificationTime, targetDate, internalCallContext);
+                                                                   final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final DateAndTimeZoneContext accountDateAndTimeZoneContext, final InternalCallContext internalCallContext) {
+        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, accountId, subscriptionId, Boolean.TRUE, futureNotificationTime, targetDate, accountDateAndTimeZoneContext, internalCallContext);
     }
 
 
     private void insertNextBillingFromTransactionInternal(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                                   final UUID subscriptionId, final Boolean isDryRunForInvoiceNotification,  final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext) {
+                                                                   final UUID subscriptionId, final Boolean isDryRunForInvoiceNotification,  final DateTime futureNotificationTime, final DateTime targetDate, final DateAndTimeZoneContext accountDateAndTimeZoneContext, final InternalCallContext internalCallContext) {
         final NotificationQueue nextBillingQueue;
         try {
             nextBillingQueue = notificationQueueService.getNotificationQueue(DefaultInvoiceService.INVOICE_SERVICE_NAME,
                                                                              DefaultNextBillingDateNotifier.NEXT_BILLING_DATE_NOTIFIER_QUEUE);
-            log.info("Queuing next billing date notification at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
 
-            nextBillingQueue.recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory.getHandle().getConnection(), futureNotificationTime,
-                                                                     new NextBillingDateNotificationKey(subscriptionId, targetDate, isDryRunForInvoiceNotification), internalCallContext.getUserToken(),
-                                                                     internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
+            // If we see existing notification for the same date (and isDryRunForInvoiceNotification mode), we don't insert a new notification
+            final List<NotificationEventWithMetadata<NextBillingDateNotificationKey>> futureNotifications =  nextBillingQueue.getFutureNotificationFromTransactionForSearchKeys(internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId(), entitySqlDaoWrapperFactory.getHandle().getConnection());
+            final NotificationEventWithMetadata<NextBillingDateNotificationKey> existingFutureNotificationWithSameDate = Iterables.tryFind(futureNotifications, new Predicate<NotificationEventWithMetadata<NextBillingDateNotificationKey>>() {
+                @Override
+                public boolean apply(final NotificationEventWithMetadata<NextBillingDateNotificationKey> input) {
+                    final boolean isEventDryRunForNotifications = input.getEvent().isDryRunForInvoiceNotification() != null ?
+                                                                  input.getEvent().isDryRunForInvoiceNotification() : false;
+
+                    final LocalDate notificationEffectiveLocaleDate = new LocalDate(futureNotificationTime, accountDateAndTimeZoneContext.getAccountTimeZone());
+                    final LocalDate eventEffectiveLocaleDate = new LocalDate(input.getEffectiveDate(), accountDateAndTimeZoneContext.getAccountTimeZone());
+
+                    return notificationEffectiveLocaleDate.compareTo(eventEffectiveLocaleDate) == 0 &&
+                           ((isDryRunForInvoiceNotification && isEventDryRunForNotifications) ||
+                            (!isDryRunForInvoiceNotification && !isEventDryRunForNotifications));
+                }
+            }).orNull();
+
+            if (existingFutureNotificationWithSameDate == null) {
+                log.info("Queuing next billing date notification at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
+
+                nextBillingQueue.recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory.getHandle().getConnection(), futureNotificationTime,
+                                                                         new NextBillingDateNotificationKey(subscriptionId, targetDate, isDryRunForInvoiceNotification), internalCallContext.getUserToken(),
+                                                                         internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
+            } else if (log.isDebugEnabled()) {
+                    log.debug("*********************   SKIPPING Queuing next billing date notification at {} for subscriptionId {} *******************", futureNotificationTime.toString(), subscriptionId.toString());
+            }
+
         } catch (final NoSuchNotificationQueue e) {
             log.error("Attempting to put items on a non-existent queue (NextBillingDateNotifier).", e);
         } catch (final IOException e) {
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java b/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
index b702434..233191e 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
@@ -23,13 +23,14 @@ import java.util.UUID;
 import org.joda.time.DateTime;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.util.entity.dao.EntitySqlDaoWrapperFactory;
+import org.killbill.billing.util.timezone.DateAndTimeZoneContext;
 
 public interface NextBillingDatePoster {
 
     void insertNextBillingNotificationFromTransaction(EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, UUID accountId,
-                                                      UUID subscriptionId, DateTime futureNotificationTime, InternalCallContext internalCallContext);
+                                                      UUID subscriptionId, DateTime futureNotificationTime, final DateAndTimeZoneContext accountDateAndTimeZoneContext, InternalCallContext internalCallContext);
 
     void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                                   final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext);
+                                                            final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final DateAndTimeZoneContext accountDateAndTimeZoneContext, final InternalCallContext internalCallContext);
 
 }
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java b/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
index cf40438..21034f5 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/dao/MockInvoiceDao.java
@@ -32,6 +32,7 @@ import org.joda.time.LocalDate;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
 import org.killbill.billing.invoice.api.user.DefaultInvoiceCreationEvent;
@@ -60,7 +61,7 @@ public class MockInvoiceDao extends MockEntityDaoBase<InvoiceModelDao, Invoice, 
 
     @Override
     public void createInvoice(final InvoiceModelDao invoice, final List<InvoiceItemModelDao> invoiceItems,
-                              final boolean isRealInvoice, final Map<UUID, List<DateTime>> callbackDateTimePerSubscriptions, final InternalCallContext context) {
+                              final boolean isRealInvoice, final FutureAccountNotifications callbackDateTimePerSubscriptions, final InternalCallContext context) {
         synchronized (monitor) {
             storeInvoice(invoice, context);
         }
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceDispatcher.java b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceDispatcher.java
index edb4c64..521ea31 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceDispatcher.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceDispatcher.java
@@ -39,6 +39,7 @@ import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.catalog.api.PhaseType;
 import org.killbill.billing.catalog.api.Plan;
 import org.killbill.billing.catalog.api.PlanPhase;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
 import org.killbill.billing.invoice.TestInvoiceHelper.DryRunFutureDateArguments;
 import org.killbill.billing.invoice.api.DryRunArguments;
 import org.killbill.billing.invoice.api.Invoice;
@@ -208,11 +209,11 @@ public class TestInvoiceDispatcher extends InvoiceTestSuiteWithEmbeddedDB {
                                                                    internalCallContextFactory, invoiceNotifier, invoicePluginDispatcher, locker, busService.getBus(),
                                                                    clock);
 
-        final Map<UUID, List<DateTime>> result = dispatcher.createNextFutureNotificationDate(Collections.singletonList(item), null, dateAndTimeZoneContext);
+        final FutureAccountNotifications futureAccountNotifications = dispatcher.createNextFutureNotificationDate(Collections.singletonList(item), null, dateAndTimeZoneContext);
 
-        Assert.assertEquals(result.size(), 1);
+        Assert.assertEquals(futureAccountNotifications.getNotifications().size(), 1);
 
-        final List<DateTime> receivedDates = result.get(item.getSubscriptionId());
+        final List<DateTime> receivedDates = futureAccountNotifications.getNotifications().get(item.getSubscriptionId());
         Assert.assertEquals(receivedDates.size(), 1);
 
         final LocalDate receivedTargetDate = new LocalDate(receivedDates.get(0), DateTimeZone.forID("Pacific/Pitcairn"));
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 5086d46..2cd3410 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/TestInvoiceHelper.java
@@ -49,6 +49,7 @@ 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.api.DryRunArguments;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
@@ -76,6 +77,7 @@ import org.killbill.billing.subscription.api.user.SubscriptionBaseApiException;
 import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
 import org.killbill.billing.util.currency.KillBillMoney;
+import org.killbill.billing.util.timezone.DateAndTimeZoneContext;
 import org.killbill.clock.Clock;
 import org.killbill.commons.locker.GlobalLocker;
 import org.mockito.Mockito;
@@ -278,7 +280,8 @@ public class TestInvoiceHelper {
                                                                                                                                          }));
 
         // The test does not use the invoice callback notifier hence the empty map
-        invoiceDao.createInvoice(invoiceModelDao, invoiceItemModelDaos, isRealInvoiceWithItems, ImmutableMap.<UUID, List<DateTime>>of(), internalCallContext);
+        final DateAndTimeZoneContext dateAndTimeZoneContext = new DateAndTimeZoneContext(clock.getUTCNow(), DateTimeZone.UTC, clock);
+        invoiceDao.createInvoice(invoiceModelDao, invoiceItemModelDaos, isRealInvoiceWithItems, new FutureAccountNotifications(dateAndTimeZoneContext, ImmutableMap.<UUID, List<DateTime>>of()), internalCallContext);
     }
 
     public void createPayment(final InvoicePayment invoicePayment, final InternalCallContext internalCallContext) {
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/tests/InvoiceTestUtils.java b/invoice/src/test/java/org/killbill/billing/invoice/tests/InvoiceTestUtils.java
index c66af70..871915a 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/tests/InvoiceTestUtils.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/tests/InvoiceTestUtils.java
@@ -24,11 +24,13 @@ import java.util.List;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.account.api.AccountApiException;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.entity.EntityPersistenceException;
+import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
 import org.killbill.billing.invoice.TestInvoiceHelper;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
@@ -40,6 +42,7 @@ import org.killbill.billing.invoice.dao.InvoiceDao;
 import org.killbill.billing.invoice.dao.InvoiceItemModelDao;
 import org.killbill.billing.invoice.dao.InvoiceModelDao;
 import org.killbill.billing.invoice.model.FixedPriceInvoiceItem;
+import org.killbill.billing.util.timezone.DateAndTimeZoneContext;
 import org.killbill.clock.Clock;
 import org.mockito.Mockito;
 import org.testng.Assert;
@@ -103,7 +106,9 @@ public class InvoiceTestUtils {
         }
         Mockito.when(invoice.getInvoiceItems()).thenReturn(invoiceItems);
 
-        invoiceDao.createInvoice(new InvoiceModelDao(invoice), invoiceModelItems, true, ImmutableMap.<UUID, List<DateTime>>of(), internalCallContext);
+        final DateAndTimeZoneContext dateAndTimeZoneContext = new DateAndTimeZoneContext(clock.getUTCNow(), DateTimeZone.UTC, clock);
+
+        invoiceDao.createInvoice(new InvoiceModelDao(invoice), invoiceModelItems, true, new FutureAccountNotifications(dateAndTimeZoneContext, ImmutableMap.<UUID, List<DateTime>>of()), internalCallContext);
 
         return invoice;
     }
diff --git a/util/src/main/java/org/killbill/billing/util/timezone/DateAndTimeZoneContext.java b/util/src/main/java/org/killbill/billing/util/timezone/DateAndTimeZoneContext.java
index 54fec5d..5781c36 100644
--- a/util/src/main/java/org/killbill/billing/util/timezone/DateAndTimeZoneContext.java
+++ b/util/src/main/java/org/killbill/billing/util/timezone/DateAndTimeZoneContext.java
@@ -77,4 +77,7 @@ public final class DateAndTimeZoneContext {
         return computeUTCDateTimeFromLocalDate(now);
     }
 
+    public DateTimeZone getAccountTimeZone() {
+        return accountTimeZone;
+    }
 }