invoice: Fix issue with removal of invoice notifications duplicate to…

… keep track of each subscriptions . See #774

The invoice notification will now be updated and keep a list of subscriptionIds, allowing dryRun filtering logic to work as expected.
The code will rely on the new NotificationQueue#updateFutureNotification api. See killbill/killbill-commons#91fa0b10029deb8ba41542a4f79ced60ee0e279a

invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java

@@ -282,6 +282,8 @@ private Invoice processAccountWithLock(final boolean parkedAccount,
                 return null;
             }
 
+            // Avoid pulling all invoices when AUTO_INVOICING_OFF is set since we will disable invoicing later
+            // (Note that we can't return right away as we send a NullInvoice event)
             final List<Invoice> existingInvoices = billingEvents.isAccountAutoInvoiceOff() ?
                                                    ImmutableList.<Invoice>of() :
                                                    ImmutableList.<Invoice>copyOf(Collections2.transform(invoiceDao.getInvoicesByAccount(context),
@@ -312,7 +314,9 @@ public Invoice apply(final InvoiceModelDao input) {
                 }  else /* DryRunType.TARGET_DATE */ {
                     invoice = processDryRun_TARGET_DATE_Invoice(accountId, inputTargetDate, candidateTargetDates, billingEvents, existingInvoices, context);
                 }
-                filterInvoiceItemsForDryRun(filteredSubscriptionIdsForDryRun, invoice);
+                if (invoice != null) {
+                    filterInvoiceItemsForDryRun(filteredSubscriptionIdsForDryRun, invoice);
+                }
             }
             return invoice;
         } catch (final CatalogApiException e) {
@@ -736,16 +740,26 @@ private Iterable<DateTime> getNextScheduledInvoiceEffectiveDate(final Iterable<U
 
             final Collection<DateTime> effectiveDates = new LinkedList<DateTime>();
             for (final NotificationEventWithMetadata<NextBillingDateNotificationKey> input : futureNotifications) {
-                final boolean isEventForSubscription = !filteredSubscriptionIds.iterator().hasNext() || Iterables.contains(filteredSubscriptionIds, input.getEvent().getUuidKey());
+
+                // If we don't specify a filter list of subscriptionIds, we look at all events.
+                boolean isEventForSubscription = !filteredSubscriptionIds.iterator().hasNext();
+                // If we specify a filter, we keep the date if at least one of the subscriptions from the event list matches one of the subscription from our filter list
+                if (filteredSubscriptionIds.iterator().hasNext()) {
+                    for (final UUID curSubscriptionId : filteredSubscriptionIds) {
+                        if (Iterables.contains(input.getEvent().getUuidKeys(), curSubscriptionId)) {
+                            isEventForSubscription = true;
+                            break;
+                        }
+                    }
+                }
+
 
                 final boolean isEventDryRunForNotifications = input.getEvent().isDryRunForInvoiceNotification() != null ?
                                                               input.getEvent().isDryRunForInvoiceNotification() : false;
                 if (isEventForSubscription && !isEventDryRunForNotifications) {
                     effectiveDates.add(input.getEffectiveDate());
                 }
-
             }
-
             return effectiveDates;
         } catch (final NoSuchNotificationQueue noSuchNotificationQueue) {
             throw new IllegalStateException(noSuchNotificationQueue);

invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDateNotifier.java

@@ -76,20 +76,22 @@ public void handleReadyNotification(final NotificationEvent notificationKey, fin
 
                 // Just to ensure compatibility with json that might not have that targetDate field (old versions < 0.13.6)
                 final DateTime targetDate = key.getTargetDate() != null ? key.getTargetDate() : eventDate;
+                final UUID firstSubscriptionId = key.getUuidKeys().iterator().next();
                 try {
-                    final SubscriptionBase subscription = subscriptionApi.getSubscriptionFromId(key.getUuidKey(), callContextFactory.createInternalTenantContext(tenantRecordId, accountRecordId));
+
+                    final SubscriptionBase subscription = subscriptionApi.getSubscriptionFromId(firstSubscriptionId, callContextFactory.createInternalTenantContext(tenantRecordId, accountRecordId));
                     if (subscription == null) {
-                        log.warn("Unable to retrieve subscriptionId='{}' for event {}", key.getUuidKey(), key);
+                        log.warn("Unable to retrieve subscriptionId='{}' for event {}", firstSubscriptionId, key);
                         return;
                     }
                     if (key.isDryRunForInvoiceNotification() != null && // Just to ensure compatibility with json that might not have that field (old versions < 0.13.6)
                         key.isDryRunForInvoiceNotification()) {
-                        processEventForInvoiceNotification(key.getUuidKey(), targetDate, userToken, accountRecordId, tenantRecordId);
+                        processEventForInvoiceNotification(firstSubscriptionId, targetDate, userToken, accountRecordId, tenantRecordId);
                     } else {
-                        processEventForInvoiceGeneration(key.getUuidKey(), targetDate, userToken, accountRecordId, tenantRecordId);
+                        processEventForInvoiceGeneration(firstSubscriptionId, targetDate, userToken, accountRecordId, tenantRecordId);
                     }
                 } catch (SubscriptionBaseApiException e) {
-                    log.warn("Error retrieving subscriptionId='{}'", key.getUuidKey(), e);
+                    log.warn("Error retrieving subscriptionId='{}'", firstSubscriptionId, e);
                 }
             }
         };

invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java

@@ -33,6 +33,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 
 public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
@@ -47,20 +48,29 @@ public DefaultNextBillingDatePoster(final NotificationQueueService notificationQ
     }
 
     @Override
-    public void insertNextBillingNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                             final UUID subscriptionId, final DateTime futureNotificationTime, final InternalCallContext internalCallContext) {
+    public void insertNextBillingNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
+                                                             final UUID accountId,
+                                                             final UUID subscriptionId,
+                                                             final DateTime futureNotificationTime,
+                                                             final InternalCallContext internalCallContext) {
         insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionId, Boolean.FALSE, futureNotificationTime, futureNotificationTime, internalCallContext);
     }
 
     @Override
-    public void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                                   final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext) {
+    public void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
+                                                                   final UUID accountId,
+                                                                   final UUID subscriptionId,
+                                                                   final DateTime futureNotificationTime,
+                                                                   final DateTime targetDate,
+                                                                   final InternalCallContext internalCallContext) {
         insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionId, Boolean.TRUE, futureNotificationTime, targetDate, internalCallContext);
     }
 
     private void insertNextBillingFromTransactionInternal(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
-                                                          final UUID subscriptionId, final Boolean isDryRunForInvoiceNotification,
-                                                          final DateTime futureNotificationTime, final DateTime targetDate,
+                                                          final UUID subscriptionId,
+                                                          final Boolean isDryRunForInvoiceNotification,
+                                                          final DateTime futureNotificationTime,
+                                                          final DateTime targetDate,
                                                           final InternalCallContext internalCallContext) {
         final NotificationQueue nextBillingQueue;
         try {
@@ -70,7 +80,7 @@ private void insertNextBillingFromTransactionInternal(final EntitySqlDaoWrapperF
             // If we see existing notification for the same date (and isDryRunForInvoiceNotification mode), we don't insert a new notification
             final Iterable<NotificationEventWithMetadata<NextBillingDateNotificationKey>> futureNotifications = nextBillingQueue.getFutureNotificationFromTransactionForSearchKeys(internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId(), entitySqlDaoWrapperFactory.getHandle().getConnection());
 
-            boolean existingFutureNotificationWithSameDate = false;
+            NotificationEventWithMetadata<NextBillingDateNotificationKey> existingNotificationForEffectiveDate = null;
             for (final NotificationEventWithMetadata<NextBillingDateNotificationKey> input : futureNotifications) {
                 final boolean isEventDryRunForNotifications = input.getEvent().isDryRunForInvoiceNotification() != null ?
                                                               input.getEvent().isDryRunForInvoiceNotification() : false;
@@ -81,19 +91,22 @@ private void insertNextBillingFromTransactionInternal(final EntitySqlDaoWrapperF
                 if (notificationEffectiveLocaleDate.compareTo(eventEffectiveLocaleDate) == 0 &&
                     ((isDryRunForInvoiceNotification && isEventDryRunForNotifications) ||
                      (!isDryRunForInvoiceNotification && !isEventDryRunForNotifications))) {
-                    existingFutureNotificationWithSameDate = true;
+                    existingNotificationForEffectiveDate = input;
                 }
                 // Go through all results to close the connection
             }
 
-            if (!existingFutureNotificationWithSameDate) {
+            if (existingNotificationForEffectiveDate == null) {
                 log.info("Queuing next billing date notification at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
 
+                final NextBillingDateNotificationKey newNotificationEvent = new NextBillingDateNotificationKey(null, ImmutableList.<UUID>of(subscriptionId), targetDate, isDryRunForInvoiceNotification);
                 nextBillingQueue.recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory.getHandle().getConnection(), futureNotificationTime,
-                                                                         new NextBillingDateNotificationKey(subscriptionId, targetDate, isDryRunForInvoiceNotification), internalCallContext.getUserToken(),
+                                                                         newNotificationEvent, 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());
+                log.info("Updating next billing date notification event at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
+                final NextBillingDateNotificationKey updateNotificationEvent = new NextBillingDateNotificationKey(existingNotificationForEffectiveDate.getEvent(), ImmutableList.of(subscriptionId));
+                nextBillingQueue.updateFutureNotification(existingNotificationForEffectiveDate.getRecordId(), updateNotificationEvent, internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
             }
 
         } catch (final NoSuchNotificationQueue e) {

invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDateNotificationKey.java

@@ -23,21 +23,35 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 
 public class NextBillingDateNotificationKey extends DefaultUUIDNotificationKey {
 
     private Boolean isDryRunForInvoiceNotification;
     private DateTime targetDate;
+    private final Iterable<UUID> uuidKeys;
 
     @JsonCreator
-    public NextBillingDateNotificationKey(@JsonProperty("uuidKey") final UUID uuidKey,
+    public NextBillingDateNotificationKey(@Deprecated @JsonProperty("uuidKey") final UUID uuidKey,
+                                          @JsonProperty("uuidKeys") final Iterable<UUID> uuidKeys,
                                           @JsonProperty("targetDate") final DateTime targetDate,
                                           @JsonProperty("isDryRunForInvoiceNotification") final Boolean isDryRunForInvoiceNotification) {
         super(uuidKey);
+        this.uuidKeys = uuidKeys;
         this.targetDate = targetDate;
         this.isDryRunForInvoiceNotification = isDryRunForInvoiceNotification;
     }
 
+    public NextBillingDateNotificationKey(NextBillingDateNotificationKey existing,
+                                          final Iterable<UUID> newUUIDKeys){
+        super(null);
+        this.uuidKeys = ImmutableList.copyOf(Iterables.concat(existing.getUuidKeys(), newUUIDKeys));
+        this.targetDate = existing.getTargetDate();
+        this.isDryRunForInvoiceNotification = existing.isDryRunForInvoiceNotification();
+    }
+
+
     @JsonProperty("isDryRunForInvoiceNotification")
     public Boolean isDryRunForInvoiceNotification() {
         return isDryRunForInvoiceNotification;
@@ -46,4 +60,13 @@ public Boolean isDryRunForInvoiceNotification() {
     public DateTime getTargetDate() {
         return targetDate;
     }
+
+    public final Iterable<UUID> getUuidKeys() {
+        // Deprecated mode
+        if (uuidKeys == null || !uuidKeys.iterator().hasNext()) {
+            return ImmutableList.of(getUuidKey());
+        } else {
+            return uuidKeys;
+        }
+    }
 }

invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotificationKey.java

@@ -24,18 +24,21 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+
 public class TestNextBillingDateNotificationKey {
 
     private static final ObjectMapper mapper = new ObjectMapper();
 
     @Test(groups = "fast")
-    public void testBasic() throws Exception {
+    public void testBasicWithUUIDKey() throws Exception {
 
         final UUID uuidKey = UUID.randomUUID();
         final DateTime targetDate = new DateTime();
         final Boolean isDryRunForInvoiceNotification = Boolean.FALSE;
 
-        final NextBillingDateNotificationKey key = new NextBillingDateNotificationKey(uuidKey, targetDate, isDryRunForInvoiceNotification);
+        final NextBillingDateNotificationKey key = new NextBillingDateNotificationKey(uuidKey, null, targetDate, isDryRunForInvoiceNotification);
         final String json = mapper.writeValueAsString(key);
 
         final NextBillingDateNotificationKey result = mapper.readValue(json, NextBillingDateNotificationKey.class);
@@ -44,6 +47,28 @@ public void testBasic() throws Exception {
         Assert.assertEquals(result.isDryRunForInvoiceNotification(), isDryRunForInvoiceNotification);
     }
 
+
+    @Test(groups = "fast")
+    public void testBasicWithUUIDKeys() throws Exception {
+
+        final UUID uuidKey1 = UUID.randomUUID();
+        final UUID uuidKey2 = UUID.randomUUID();
+        final DateTime targetDate = new DateTime();
+        final Boolean isDryRunForInvoiceNotification = Boolean.FALSE;
+
+        final NextBillingDateNotificationKey key = new NextBillingDateNotificationKey(null, ImmutableList.of(uuidKey1, uuidKey2), targetDate, isDryRunForInvoiceNotification);
+        final String json = mapper.writeValueAsString(key);
+
+        final NextBillingDateNotificationKey result = mapper.readValue(json, NextBillingDateNotificationKey.class);
+        Assert.assertNull(result.getUuidKey());
+        Assert.assertEquals(result.getTargetDate().compareTo(targetDate), 0);
+        Assert.assertEquals(result.isDryRunForInvoiceNotification(), isDryRunForInvoiceNotification);
+        Assert.assertNotNull(result.getUuidKeys());
+
+        Assert.assertTrue(Iterables.contains(result.getUuidKeys(), uuidKey1));
+        Assert.assertTrue(Iterables.contains(result.getUuidKeys(), uuidKey2));
+    }
+
     @Test(groups = "fast")
     public void testWithMissingFields() throws Exception {
         final String json = "{\"uuidKey\":\"a38c363f-b25b-4287-8ebc-55964e116d2f\"}";
@@ -52,5 +77,9 @@ public void testWithMissingFields() throws Exception {
         Assert.assertNull(result.getTargetDate());
         Assert.assertNull(result.isDryRunForInvoiceNotification());
 
+        // Compatibility mode : Although the  uuidKeys is not in the json, we verify the getter return the right result
+        Assert.assertNotNull(result.getUuidKeys());
+        Assert.assertEquals(result.getUuidKeys().iterator().next().toString(), "a38c363f-b25b-4287-8ebc-55964e116d2f");
+
     }
 }

invoice/src/test/java/org/killbill/billing/invoice/notification/TestNextBillingDateNotifier.java

@@ -31,6 +31,8 @@
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import com.google.common.collect.ImmutableList;
+
 import static org.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.MINUTES;
 
@@ -47,7 +49,7 @@ public void testInvoiceNotifier() throws Exception {
 
         final NotificationQueue nextBillingQueue = notificationQueueService.getNotificationQueue(DefaultInvoiceService.INVOICE_SERVICE_NAME, DefaultNextBillingDateNotifier.NEXT_BILLING_DATE_NOTIFIER_QUEUE);
 
-        nextBillingQueue.recordFutureNotification(now, new NextBillingDateNotificationKey(subscriptionId, now, Boolean.FALSE), internalCallContext.getUserToken(), accountRecordId, internalCallContext.getTenantRecordId());
+        nextBillingQueue.recordFutureNotification(now, new NextBillingDateNotificationKey(null, ImmutableList.<UUID>of(subscriptionId), now, Boolean.FALSE), internalCallContext.getUserToken(), accountRecordId, internalCallContext.getTenantRecordId());
 
         // Move time in the future after the notification effectiveDate
         clock.setDeltaFromReality(3000);