killbill-aplcache

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 9139a3e..7d74b2d 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
@@ -17,6 +17,7 @@
 package com.ning.billing.beatrix.integration;
 
 import java.math.BigDecimal;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.UUID;
 
@@ -36,6 +37,9 @@ import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.entitlement.api.user.SubscriptionBundle;
 import com.ning.billing.entitlement.api.user.SubscriptionData;
 import com.ning.billing.invoice.api.Invoice;
+import com.ning.billing.invoice.generator.InvoiceDateUtils;
+
+import com.google.common.collect.ImmutableList;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
@@ -45,30 +49,45 @@ import static org.testng.Assert.assertTrue;
 public class TestIntegration extends TestIntegrationBase {
     @Test(groups = "slow")
     public void testBasePlanCompleteWithBillingDayInPast() throws Exception {
-        log.info("Starting testBasePlanCompleteWithBillingDayInPast");
         final DateTime startDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
-        testBasePlanComplete(startDate, 31, false);
-    }
-
-    @Test(groups = "slow")
-    public void testBasePlanCompleteWithBillingDayPresent() throws Exception {
-        log.info("Starting testBasePlanCompleteWithBillingDayPresent");
-        final DateTime startDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
-        testBasePlanComplete(startDate, 1, false);
+        final LinkedHashMap<DateTime, List<NextEvent>> expectedStates = new LinkedHashMap<DateTime, List<NextEvent>>();
+        expectedStates.put(new DateTime(2012, 2, 28, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 2, 29, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 1, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 2, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of(NextEvent.PHASE,
+                                                                                                            NextEvent.INVOICE,
+                                                                                                            NextEvent.PAYMENT));
+        expectedStates.put(new DateTime(2012, 3, 3, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        testBasePlanComplete(startDate, expectedStates, 31, false);
     }
 
     @Test(groups = "slow")
     public void testBasePlanCompleteWithBillingDayAlignedWithTrial() throws Exception {
-        log.info("Starting testBasePlanCompleteWithBillingDayAlignedWithTrial");
         final DateTime startDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
-        testBasePlanComplete(startDate, 2, false);
+        final LinkedHashMap<DateTime, List<NextEvent>> expectedStates = new LinkedHashMap<DateTime, List<NextEvent>>();
+        expectedStates.put(new DateTime(2012, 2, 28, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 2, 29, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 1, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 2, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of(NextEvent.PHASE,
+                                                                                                            NextEvent.INVOICE,
+                                                                                                            NextEvent.PAYMENT));
+        expectedStates.put(new DateTime(2012, 3, 3, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        testBasePlanComplete(startDate, expectedStates, 2, false);
     }
 
     @Test(groups = "slow")
     public void testBasePlanCompleteWithBillingDayInFuture() throws Exception {
-        log.info("Starting testBasePlanCompleteWithBillingDayInFuture");
         final DateTime startDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
-        testBasePlanComplete(startDate, 3, true);
+        final LinkedHashMap<DateTime, List<NextEvent>> expectedStates = new LinkedHashMap<DateTime, List<NextEvent>>();
+        expectedStates.put(new DateTime(2012, 2, 28, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 2, 29, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 1, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of());
+        expectedStates.put(new DateTime(2012, 3, 2, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of(NextEvent.PHASE,
+                                                                                                            NextEvent.INVOICE,
+                                                                                                            NextEvent.PAYMENT));
+        expectedStates.put(new DateTime(2012, 3, 3, 0, 3, 45, 0, testTimeZone), ImmutableList.<NextEvent>of(NextEvent.INVOICE,
+                                                                                                            NextEvent.PAYMENT));
+        testBasePlanComplete(startDate, expectedStates, 3, true);
     }
 
     @Test(groups = {"stress"})
@@ -80,9 +99,6 @@ public class TestIntegration extends TestIntegrationBase {
             }
 
             log.info("################################  ITERATION " + curIteration + "  #########################");
-            Thread.sleep(1000);
-            testBasePlanCompleteWithBillingDayPresent();
-            Thread.sleep(1000);
             cleanupTest();
             setupTest();
             testBasePlanCompleteWithBillingDayInPast();
@@ -98,11 +114,9 @@ public class TestIntegration extends TestIntegrationBase {
                 cleanupTest();
                 Thread.sleep(1000);
             }
-
         }
     }
 
-
     @Test(groups = {"stress"})
     public void stressTestDebug() throws Exception {
         final int maxIterations = 100;
@@ -122,9 +136,6 @@ public class TestIntegration extends TestIntegrationBase {
 
     @Test(groups = "slow")
     public void testAddonsWithMultipleAlignments() throws Exception {
-
-        log.info("Starting testRepairChangeBPWithAddonIncluded");
-
         final DateTime initialDate = new DateTime(2012, 4, 25, 0, 13, 42, 0, testTimeZone);
         clock.setDeltaFromReality(initialDate.getMillis() - clock.getUTCNow().getMillis());
 
@@ -140,7 +151,7 @@ public class TestIntegration extends TestIntegrationBase {
         busHandler.pushExpectedEvent(NextEvent.CREATE);
         busHandler.pushExpectedEvent(NextEvent.INVOICE);
         final SubscriptionData baseSubscription = subscriptionDataFromSubscription(entitlementUserApi.createSubscription(bundle.getId(),
-                                                                                                                   new PlanPhaseSpecifier(productName, ProductCategory.BASE, term, planSetName, null), null, context));
+                                                                                                                         new PlanPhaseSpecifier(productName, ProductCategory.BASE, term, planSetName, null), null, context));
         assertNotNull(baseSubscription);
         assertTrue(busHandler.isCompleted(DELAY));
         assertListenerStatus();
@@ -163,7 +174,7 @@ public class TestIntegration extends TestIntegrationBase {
         busHandler.pushExpectedEvent(NextEvent.INVOICE);
         busHandler.pushExpectedEvent(NextEvent.PAYMENT);
         final SubscriptionData aoSubscription2 = subscriptionDataFromSubscription(entitlementUserApi.createSubscription(bundle.getId(),
-                                                                                                                  new PlanPhaseSpecifier("Laser-Scope", ProductCategory.ADD_ON, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null), null, context));
+                                                                                                                        new PlanPhaseSpecifier("Laser-Scope", ProductCategory.ADD_ON, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null), null, context));
         assertTrue(busHandler.isCompleted(DELAY));
         assertListenerStatus();
 
@@ -238,9 +249,6 @@ public class TestIntegration extends TestIntegrationBase {
 
     @Test(groups = "slow")
     public void testWithRecreatePlan() throws Exception {
-
-        log.info("Starting testWithRecreatePlan");
-
         final DateTime initialDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
         final int billingDay = 2;
 
@@ -310,8 +318,8 @@ public class TestIntegration extends TestIntegrationBase {
         assertListenerStatus();
     }
 
-    private void testBasePlanComplete(final DateTime initialCreationDate, final int billingDay,
-                                      final boolean proRationExpected) throws Exception {
+    private void testBasePlanComplete(final DateTime initialCreationDate, final LinkedHashMap<DateTime, List<NextEvent>> expectedStates,
+                                      final int billingDay, final boolean proRationExpected) throws Exception {
         log.info("Beginning test with BCD of " + billingDay);
         final Account account = createAccountWithPaymentMethod(getAccountData(billingDay));
         final UUID accountId = account.getId();
@@ -365,25 +373,16 @@ public class TestIntegration extends TestIntegrationBase {
         verifyTestResult(accountId, subscription.getId(), startDate, endDate, rate, endDate, invoiceItemCount);
 
         //
-        // MOVE TIME TO AFTER TRIAL AND EXPECT BOTH EVENTS :  NextEvent.PHASE NextEvent.INVOICE
+        // MOVE TIME
         //
-        busHandler.pushExpectedEvent(NextEvent.PHASE);
-        busHandler.pushExpectedEvent(NextEvent.INVOICE);
-        busHandler.pushExpectedEvent(NextEvent.PAYMENT);
-
-        if (proRationExpected) {
-            busHandler.pushExpectedEvent(NextEvent.INVOICE);
-            busHandler.pushExpectedEvent(NextEvent.PAYMENT);
+        for (final DateTime dateTimeToSet : expectedStates.keySet()) {
+            for (final NextEvent expectedEvent : expectedStates.get(dateTimeToSet)) {
+                busHandler.pushExpectedEvent(expectedEvent);
+            }
+            clock.setTime(dateTimeToSet);
+            assertTrue(busHandler.isCompleted(DELAY));
         }
 
-        clock.addDeltaFromReality(AT_LEAST_ONE_MONTH_MS);
-
-        // STEPH
-        /*
-        Thread.sleep(6000000);
-        */
-        assertTrue(busHandler.isCompleted(DELAY));
-
         startDate = subscription.getCurrentPhaseStart();
         rate = subscription.getCurrentPhase().getRecurringPrice().getPrice(Currency.USD);
         BigDecimal price;
@@ -440,7 +439,7 @@ public class TestIntegration extends TestIntegrationBase {
         busHandler.pushExpectedEvent(NextEvent.INVOICE);
         busHandler.pushExpectedEvent(NextEvent.PAYMENT);
         //clock.addDeltaFromReality(ctd.getMillis() - clock.getUTCNow().getMillis());
-        clock.addDeltaFromReality(AT_LEAST_ONE_MONTH_MS + 1000);
+        clock.addDeltaFromReality(AT_LEAST_ONE_MONTH_MS);
 
         //waitForDebug();
 

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 f20ac76..de351c5 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
@@ -90,7 +90,7 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
     protected static final BigDecimal THIRTY_ONE = new BigDecimal("31.0000").setScale(NUMBER_OF_DECIMALS);
 
     protected static final Logger log = LoggerFactory.getLogger(TestIntegration.class);
-    protected static long AT_LEAST_ONE_MONTH_MS = 31L * 24L * 3600L * 1000L;
+    protected static long AT_LEAST_ONE_MONTH_MS = 31L * 24L * 3600L * 1000L + 2000L;
 
     protected static final long DELAY = 5000;
 

diff --git a/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
index 4652130..86d7bef 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/dao/DefaultInvoiceDao.java
@@ -538,7 +538,7 @@ public class DefaultInvoiceDao implements InvoiceDao {
         if (shouldBeNotified) {
             // We could be notified at any time during the day at the billCycleDay - use the current time to
             // spread the load
-            final DateTime nextNotificationDateTime = InvoiceDateUtils.calculateBillingCycleDateOnOrAfter(clock.getUTCNow(), billCycleDay);
+            final DateTime nextNotificationDateTime = InvoiceDateUtils.calculateBillingCycleDateAfter(clock.getUTCNow(), billCycleDay);
             // NextBillingDatePoster will ignore duplicates
             nextBillingDatePoster.insertNextBillingNotification(dao, accountId, subscriptionForNextNotification, nextNotificationDateTime);
         }

diff --git a/invoice/src/main/java/com/ning/billing/invoice/generator/InvoiceDateUtils.java b/invoice/src/main/java/com/ning/billing/invoice/generator/InvoiceDateUtils.java
index df130ee..5b2b60f 100644
--- a/invoice/src/main/java/com/ning/billing/invoice/generator/InvoiceDateUtils.java
+++ b/invoice/src/main/java/com/ning/billing/invoice/generator/InvoiceDateUtils.java
@@ -30,6 +30,7 @@ public class InvoiceDateUtils {
         return new DateTime(tzAdjustedStartDate.getYear(), tzAdjustedStartDate.getMonthOfYear(), tzAdjustedStartDate.getDayOfMonth(), 0, 0, timeZone);
     }
 
+    // Note: date has to be in UTC
     public static DateTime calculateBillingCycleDateOnOrAfter(final DateTime date, final int billingCycleDay) {
         final int lastDayOfMonth = date.dayOfMonth().getMaximumValue();
 
@@ -46,4 +47,14 @@ public class InvoiceDateUtils {
         }
         return proposedDate;
     }
+
+    // Note: date has to be in UTC
+    public static DateTime calculateBillingCycleDateAfter(final DateTime date, final int billingCycleDay) {
+        DateTime proposedDate = calculateBillingCycleDateOnOrAfter(date, billingCycleDay);
+
+        if (date.compareTo(proposedDate) == 0) {
+            proposedDate = proposedDate.plusMonths(1);
+        }
+        return proposedDate;
+    }
 }

diff --git a/invoice/src/test/java/com/ning/billing/invoice/generator/TestInvoiceDateUtils.java b/invoice/src/test/java/com/ning/billing/invoice/generator/TestInvoiceDateUtils.java
new file mode 100644
index 0000000..810e76b
--- /dev/null
+++ b/invoice/src/test/java/com/ning/billing/invoice/generator/TestInvoiceDateUtils.java
@@ -0,0 +1,66 @@
+/*
+ * 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.invoice.generator;
+
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class TestInvoiceDateUtils {
+    @Test(groups = "fast")
+    public void testBeforeBCDWithAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-02T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-03-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+
+    @Test(groups = "fast")
+    public void testEqualBCDWithAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-03T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-04-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+
+    @Test(groups = "fast")
+    public void testAfterBCDWithAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-04T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-04-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+
+    @Test(groups = "fast")
+    public void testBeforeBCDWithOnOrAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-02T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateOnOrAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-03-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+
+    @Test(groups = "fast")
+    public void testEqualBCDWithOnOrAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-03T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateOnOrAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-03-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+
+    @Test(groups = "fast")
+    public void testAfterBCDWithOnOrAfter() throws Exception {
+        final DateTime from = new DateTime("2012-03-04T00:03:47.000Z", DateTimeZone.UTC);
+        final DateTime to = InvoiceDateUtils.calculateBillingCycleDateOnOrAfter(from, 3);
+        Assert.assertEquals(to, new DateTime("2012-04-03T00:03:47.000Z", DateTimeZone.UTC));
+    }
+}