killbill-aplcache

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoice.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoice.java
index 9e8941c..3d81d73 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoice.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoice.java
@@ -189,7 +189,6 @@ public class TestIntegrationInvoice extends TestIntegrationBase {
         // We test first the next expected invoice for a specific subscription: We can see the targetDate is 2015-2-14 and not 2015-2-1
         final DryRunArguments dryRunWIthSubscription = new TestDryRunArguments(DryRunType.UPCOMING_INVOICE, null, null, null, null, null, null, subscriptionMonthly.getId(), null, null, null);
         expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("249.95")));
-        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 1), new LocalDate(2016, 2, 1), InvoiceItemType.RECURRING, new BigDecimal("2399.95")));
         dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, dryRunWIthSubscription, callContext);
         assertEquals(dryRunInvoice.getTargetDate(), new LocalDate(2015, 2, 14));
         invoiceChecker.checkInvoiceNoAudits(dryRunInvoice, callContext, expectedInvoices);

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 7764c4a..49ec8d8 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
@@ -22,6 +22,7 @@ import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -235,12 +236,14 @@ public class InvoiceDispatcher {
             if (billingEvents.isEmpty()) {
                 return null;
             }
+            final Iterable<UUID> filteredSubscriptionIdsForDryRun = getFilteredSubscriptionIdsForDryRun(dryRunArguments, billingEvents);
             final List<DateTime> candidateDateTimes = (inputTargetDateTime != null) ?
                                                       ImmutableList.of(inputTargetDateTime) :
-                                                      getUpcomingInvoiceCandidateDates(getFilteredSubscriptionIds(dryRunArguments, billingEvents), context);
+                                                      getUpcomingInvoiceCandidateDates(filteredSubscriptionIdsForDryRun, context);
             for (final DateTime curTargetDateTime : candidateDateTimes) {
                 final Invoice invoice = processAccountWithLockAndInputTargetDate(accountId, curTargetDateTime, billingEvents, isDryRun, context);
                 if (invoice != null) {
+                    filterInvoiceItemsForDryRun(filteredSubscriptionIdsForDryRun, invoice);
                     return invoice;
                 }
             }
@@ -251,9 +254,23 @@ public class InvoiceDispatcher {
         }
     }
 
+    private void filterInvoiceItemsForDryRun(final Iterable<UUID> filteredSubscriptionIdsForDryRun, final Invoice invoice) {
+        if (!filteredSubscriptionIdsForDryRun.iterator().hasNext()) {
+            return;
+        }
+
+        final Iterator<InvoiceItem> it = invoice.getInvoiceItems().iterator();
+        while (it.hasNext()) {
+            final InvoiceItem cur = it.next();
+            if (!Iterables.contains(filteredSubscriptionIdsForDryRun, cur.getSubscriptionId())) {
+                it.remove();
+            }
+        }
+    }
 
-    private Iterable<UUID> getFilteredSubscriptionIds(final DryRunArguments dryRunArguments, final BillingEventSet billingEvents) {
-        if (!dryRunArguments.getDryRunType().equals(DryRunType.UPCOMING_INVOICE) ||
+    private Iterable<UUID> getFilteredSubscriptionIdsForDryRun(@Nullable final DryRunArguments dryRunArguments, final BillingEventSet billingEvents) {
+        if (dryRunArguments == null ||
+            !dryRunArguments.getDryRunType().equals(DryRunType.UPCOMING_INVOICE) ||
             (dryRunArguments.getSubscriptionId() == null && dryRunArguments.getBundleId() == null)) {
             return ImmutableList.<UUID>of();
         }