killbill-memoizeit

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
index 998ac96..a50b665 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/overdue/TestOverdueIntegration.java
@@ -727,7 +727,7 @@ public class TestOverdueIntegration extends TestOverdueBase {
             }
         } else {
             // Upgrade - we don't expect a payment here due to the scenario (the account will have some CBA)
-            changeSubscriptionAndCheckForCompletion(subscription, "Assault-Rifle", BillingPeriod.MONTHLY, NextEvent.CHANGE, NextEvent.INVOICE);
+            changeSubscriptionAndCheckForCompletion(subscription, "Assault-Rifle", BillingPeriod.MONTHLY, NextEvent.CHANGE, NextEvent.INVOICE, NextEvent.INVOICE_ADJUSTMENT);
         }
     }
 }

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 e871545..ff8328c 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
@@ -27,6 +27,7 @@ import javax.annotation.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
+import org.skife.jdbi.v2.IDBI;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
@@ -116,7 +117,7 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
     protected static final Logger log = LoggerFactory.getLogger(TestIntegration.class);
     protected static long AT_LEAST_ONE_MONTH_MS = 32L * 24L * 3600L * 1000L;
 
-    protected static final long DELAY = 5000;
+    protected static final long DELAY = 10000;
 
 
     @Inject
@@ -198,6 +199,9 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
     @Inject
     protected RecordIdApi recordIdApi;
 
+    @Inject
+    protected IDBI idbi;
+
     @javax.inject.Inject
     protected CacheControllerDispatcher controlCacheDispatcher;
 
@@ -229,7 +233,7 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
     public void beforeClass() throws Exception {
         final Injector g = Guice.createInjector(Stage.PRODUCTION, new BeatrixIntegrationModule(configSource));
         g.injectMembers(this);
-        busHandler = new TestApiListener(this);
+        busHandler = new TestApiListener(this, idbi);
 
         SetupBundleWithAssertion setupTest = new SetupBundleWithAssertion("whatever", osgiConfig, "whatever");
         setupTest.cleanBundleInstallDir();

diff --git a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
index aaca80d..ef9d1b2 100644
--- a/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
+++ b/beatrix/src/test/java/com/ning/billing/beatrix/integration/TestIntegrationWithAutoInvoiceOffTag.java
@@ -78,6 +78,7 @@ public class TestIntegrationWithAutoInvoiceOffTag extends TestIntegrationBase {
     @Test(groups = {"slow"}, enabled = true)
     public void testAutoInvoiceOffAccount() throws Exception {
         clock.setTime(new DateTime(2012, 5, 1, 0, 3, 42, 0));
+        busHandler.pushExpectedEvents(NextEvent.TAG);
         add_AUTO_INVOICING_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
 
         // set next invoice to fail and create network
@@ -104,7 +105,7 @@ public class TestIntegrationWithAutoInvoiceOffTag extends TestIntegrationBase {
         invoices = invoiceApi.getInvoicesByAccount(account.getId(), callContext);
         assertEquals(invoices.size(), 0);
 
-        busHandler.pushExpectedEvents(NextEvent.INVOICE);
+        busHandler.pushExpectedEvents(NextEvent.TAG, NextEvent.INVOICE, NextEvent.PAYMENT);
         remove_AUTO_INVOICING_OFF_Tag(account.getId(), ObjectType.ACCOUNT);
         assertTrue(busHandler.isCompleted(DELAY));
 
@@ -163,7 +164,7 @@ public class TestIntegrationWithAutoInvoiceOffTag extends TestIntegrationBase {
 
         add_AUTO_INVOICING_OFF_Tag(baseSubscription.getBundleId(), ObjectType.BUNDLE);
 
-        busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.PHASE, NextEvent.INVOICE);
+        busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.PHASE, NextEvent.INVOICE, NextEvent.PAYMENT);
         clock.addDays(40); // DAY 40 out of trial
         assertTrue(busHandler.isCompleted(DELAY));
 

diff --git a/server/src/test/java/com/ning/billing/jaxrs/TestJaxrsBase.java b/server/src/test/java/com/ning/billing/jaxrs/TestJaxrsBase.java
index dc5186f..7974fa3 100644
--- a/server/src/test/java/com/ning/billing/jaxrs/TestJaxrsBase.java
+++ b/server/src/test/java/com/ning/billing/jaxrs/TestJaxrsBase.java
@@ -253,7 +253,7 @@ public class TestJaxrsBase extends KillbillClient {
 
         //mapper.setPropertyNamingStrategy(new PropertyNamingStrategy.LowerCaseWithUnderscoresStrategy());
 
-        busHandler = new TestApiListener(null);
+        busHandler = new TestApiListener(null, getDBTestingHelper().getDBI());
     }
 
     protected void loadConfig() {

diff --git a/util/src/test/java/com/ning/billing/api/TestApiListener.java b/util/src/test/java/com/ning/billing/api/TestApiListener.java
index e01ff24..b2319ba 100644
--- a/util/src/test/java/com/ning/billing/api/TestApiListener.java
+++ b/util/src/test/java/com/ning/billing/api/TestApiListener.java
@@ -19,6 +19,7 @@ package com.ning.billing.api;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Stack;
+import java.util.concurrent.Callable;
 
 import javax.inject.Inject;
 
@@ -33,12 +34,18 @@ import com.ning.billing.util.events.RepairEntitlementInternalEvent;
 import com.ning.billing.util.events.TagDefinitionInternalEvent;
 import com.ning.billing.util.events.TagInternalEvent;
 import org.joda.time.DateTime;
+import org.skife.jdbi.v2.Handle;
+import org.skife.jdbi.v2.IDBI;
+import org.skife.jdbi.v2.tweak.HandleCallback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.google.common.base.Joiner;
 import com.google.common.eventbus.Subscribe;
 
+import static com.jayway.awaitility.Awaitility.await;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
 public class TestApiListener {
 
     protected static final Logger log = LoggerFactory.getLogger(TestApiListener.class);
@@ -46,17 +53,19 @@ public class TestApiListener {
     private final List<NextEvent> nextExpectedEvent;
 
     private final TestListenerStatus testStatus;
+    private final IDBI idbi;
 
     private boolean nonExpectedMode;
 
     private volatile boolean completed;
 
     @Inject
-    public TestApiListener(final TestListenerStatus testStatus) {
+    public TestApiListener(final TestListenerStatus testStatus, final IDBI idbi) {
         nextExpectedEvent = new Stack<NextEvent>();
         this.completed = false;
         this.testStatus = testStatus;
         this.nonExpectedMode = false;
+        this.idbi = idbi;
     }
 
     public enum NextEvent {
@@ -231,6 +240,24 @@ public class TestApiListener {
                     final DateTime before = new DateTime();
                     wait(500);
                     if (completed) {
+                        // TODO PIERRE Kludge alert!
+                        // When we arrive here, we got notified by the current thread (Bus listener) that we received
+                        // all expected events. But other handlers might still be processing them.
+                        // Since there is only one bus thread, and that the test thread waits for all events to be processed,
+                        // we're guaranteed that all are processed when the bus events table is empty.
+                        await().atMost(10, SECONDS).until(new Callable<Boolean>() {
+                            @Override
+                            public Boolean call() throws Exception {
+                                final long inProcessingBusEvents = idbi.withHandle(new HandleCallback<Long>() {
+                                    @Override
+                                    public Long withHandle(final Handle handle) throws Exception {
+                                        return (Long) handle.select("select count(distinct record_id) count from bus_events").get(0).get("count");
+                                    }
+                                });
+                                log.debug("Events still in processing: " + inProcessingBusEvents);
+                                return inProcessingBusEvents == 0;
+                            }
+                        });
                         return completed;
                     }
                     final DateTime after = new DateTime();

diff --git a/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDao.java b/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDao.java
index 90a84a4..e9ebb64 100644
--- a/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDao.java
+++ b/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDao.java
@@ -44,7 +44,7 @@ public class TestDefaultTagDao extends UtilTestSuiteWithEmbeddedDB {
     @BeforeMethod(groups = "slow")
     public void beforeMethod() throws Exception {
         super.beforeMethod();
-        eventsListener = new TestApiListener(null);
+        eventsListener = new TestApiListener(null, idbi);
         eventBus.register(eventsListener);
     }
 

diff --git a/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDefinitionDao.java b/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDefinitionDao.java
index d750dd1..908de4f 100644
--- a/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDefinitionDao.java
+++ b/util/src/test/java/com/ning/billing/util/tag/dao/TestDefaultTagDefinitionDao.java
@@ -42,7 +42,7 @@ public class TestDefaultTagDefinitionDao extends UtilTestSuiteWithEmbeddedDB {
     @BeforeMethod(groups = "slow")
     public void beforeMethod() throws Exception {
         super.beforeMethod();
-        eventsListener = new TestApiListener(null);
+        eventsListener = new TestApiListener(null, idbi);
         eventBus.register(eventsListener);
     }
 

diff --git a/util/src/test/java/com/ning/billing/util/UtilTestSuiteWithEmbeddedDB.java b/util/src/test/java/com/ning/billing/util/UtilTestSuiteWithEmbeddedDB.java
index b618f24..ee15ee9 100644
--- a/util/src/test/java/com/ning/billing/util/UtilTestSuiteWithEmbeddedDB.java
+++ b/util/src/test/java/com/ning/billing/util/UtilTestSuiteWithEmbeddedDB.java
@@ -18,6 +18,7 @@ package com.ning.billing.util;
 
 import javax.inject.Inject;
 
+import org.skife.jdbi.v2.IDBI;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
@@ -65,6 +66,8 @@ public abstract class UtilTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuite
     protected DefaultTagDao tagDao;
     @Inject
     protected AuditDao auditDao;
+    @Inject
+    protected IDBI idbi;
 
     @BeforeClass(groups = "slow")
     public void beforeClass() throws Exception {