killbill-memoizeit

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegration.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegration.java
index 89507b5..5536865 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegration.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegration.java
@@ -421,49 +421,6 @@ public class TestIntegration extends TestIntegrationBase {
         checkNoMoreInvoiceToGenerate(account);
     }
 
-    @Test(groups = {"stress"}, enabled = false)
-    public void stressTest() throws Exception {
-        final int maxIterations = 100;
-        for (int curIteration = 0; curIteration < maxIterations; curIteration++) {
-            if (curIteration != 0) {
-                beforeMethod();
-            }
-
-            log.info("################################  ITERATION " + curIteration + "  #########################");
-            afterMethod();
-            beforeMethod();
-            testBasePlanCompleteWithBillingDayInPast();
-            Thread.sleep(1000);
-            afterMethod();
-            beforeMethod();
-            testBasePlanCompleteWithBillingDayAlignedWithTrial();
-            Thread.sleep(1000);
-            afterMethod();
-            beforeMethod();
-            testBasePlanCompleteWithBillingDayInFuture();
-            if (curIteration < maxIterations - 1) {
-                afterMethod();
-                Thread.sleep(1000);
-            }
-        }
-    }
-
-    @Test(groups = {"stress"}, enabled = false)
-    public void stressTestDebug() throws Exception {
-        final int maxIterations = 100;
-        for (int curIteration = 0; curIteration < maxIterations; curIteration++) {
-            log.info("################################  ITERATION " + curIteration + "  #########################");
-            if (curIteration != 0) {
-                beforeMethod();
-            }
-            testAddonsWithMultipleAlignments();
-            if (curIteration < maxIterations - 1) {
-                afterMethod();
-                Thread.sleep(1000);
-            }
-        }
-    }
-
     @Test(groups = "slow")
     public void testAddonsWithMultipleAlignments() throws Exception {
         final DateTime initialDate = new DateTime(2012, 4, 25, 0, 13, 42, 0, testTimeZone);

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
index 2969715..1b13cf2 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
@@ -18,15 +18,18 @@
 
 package org.killbill.billing.beatrix.integration;
 
-import com.google.common.base.Function;
-import com.google.common.base.Joiner;
-import com.google.common.base.Predicate;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-import com.google.inject.Guice;
-import com.google.inject.Injector;
-import com.google.inject.Stage;
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.Callable;
+
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+import javax.inject.Named;
+
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
@@ -122,30 +125,31 @@ import org.skife.config.TimeSpan;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
+import org.testng.IHookCallBack;
+import org.testng.IHookable;
+import org.testng.ITestResult;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 
-import javax.annotation.Nullable;
-import javax.inject.Inject;
-import javax.inject.Named;
-import java.math.BigDecimal;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.UUID;
-import java.util.concurrent.Callable;
+import com.google.common.base.Function;
+import com.google.common.base.Joiner;
+import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.google.inject.Stage;
 
-import static org.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
 
-
-public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB {
+public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implements IHookable {
 
     protected static final DateTimeZone testTimeZone = DateTimeZone.UTC;
     protected static final Logger log = LoggerFactory.getLogger(TestIntegrationBase.class);
@@ -329,16 +333,10 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB {
         lifecycle.fireStartupSequencePostEventRegistration();
 
         paymentPlugin.clear();
-
-        // Make sure we start with a clean state
-        assertListenerStatus();
     }
 
     @AfterMethod(groups = "slow")
     public void afterMethod() throws Exception {
-        // Make sure we finish in a clean state
-        assertListenerStatus();
-
         lifecycle.fireShutdownSequencePriorEventUnRegistration();
         busService.getBus().unregister(busHandler);
         lifecycle.fireShutdownSequencePostEventUnRegistration();
@@ -348,6 +346,20 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB {
         log.debug("DONE WITH TEST");
     }
 
+    // Note: assertions should not be run in before / after hooks, as the associated test result won't be correctly updated.
+    // Use this wrapper instead.
+    @Override
+    public void run(final IHookCallBack callBack, final ITestResult testResult) {
+        // Make sure we start with a clean state
+        assertListenerStatus();
+
+        // Run the actual test
+        callBack.runTestMethod(testResult);
+
+        // Make sure we finish in a clean state
+        assertListenerStatus();
+    }
+
     protected void checkNoMoreInvoiceToGenerate(final Account account) {
         busHandler.pushExpectedEvent(NextEvent.NULL_INVOICE);
         try {

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
index 18f2da2..a755f9c 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
@@ -64,12 +64,6 @@ public class TestIntegrationInvoiceWithRepairLogic extends TestIntegrationBase {
     @Inject
     protected InvoiceDao invoiceDao;
 
-
-    @AfterMethod(groups = "slow")
-    public void afterMethod() throws Exception {
-        super.afterMethod();
-    }
-
     @Test(groups = "slow")
     public void testSimplePartialRepairWithItemAdjustment() throws Exception {
         // We take april as it has 30 days (easier to play with BCD)

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestPublicBus.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestPublicBus.java
index 4e59e25..c603090 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestPublicBus.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestPublicBus.java
@@ -24,13 +24,10 @@ import java.util.concurrent.Callable;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.joda.time.DateTime;
-import org.killbill.billing.DBTestingHelper;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.api.TestApiListener.NextEvent;
-import org.killbill.billing.beatrix.extbus.DefaultBusExternalEvent;
 import org.killbill.billing.callcontext.DefaultCallContext;
 import org.killbill.billing.catalog.api.BillingPeriod;
-import org.killbill.billing.catalog.api.PriceListSet;
 import org.killbill.billing.catalog.api.ProductCategory;
 import org.killbill.billing.entitlement.api.DefaultEntitlement;
 import org.killbill.billing.notification.plugin.api.ExtBusEvent;
@@ -46,11 +43,6 @@ import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.CallOrigin;
 import org.killbill.billing.util.callcontext.UserType;
 import org.killbill.billing.util.jackson.ObjectMapper;
-import org.killbill.billing.util.nodes.NodeCommand;
-import org.killbill.billing.util.nodes.NodeCommandMetadata;
-import org.killbill.billing.util.nodes.NodeCommandProperty;
-import org.killbill.billing.util.nodes.PluginNodeCommandMetadata;
-import org.killbill.billing.util.nodes.SystemNodeCommandType;
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
@@ -58,12 +50,11 @@ import org.testng.annotations.Test;
 
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.databind.JsonMappingException;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.eventbus.Subscribe;
 
-import static org.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
 import static org.testng.Assert.assertNotNull;
 
 public class TestPublicBus extends TestIntegrationBase {

diff --git a/util/src/test/java/org/killbill/billing/api/FlakyInvokedMethodListener.java b/util/src/test/java/org/killbill/billing/api/FlakyInvokedMethodListener.java
new file mode 100644
index 0000000..80ad169
--- /dev/null
+++ b/util/src/test/java/org/killbill/billing/api/FlakyInvokedMethodListener.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
+ *
+ * The Billing Project 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 org.killbill.billing.api;
+
+import org.testng.IInvokedMethod;
+import org.testng.IInvokedMethodListener;
+import org.testng.IRetryAnalyzer;
+import org.testng.ITestResult;
+import org.testng.Reporter;
+
+public class FlakyInvokedMethodListener implements IInvokedMethodListener {
+
+    @Override
+    public void beforeInvocation(final IInvokedMethod method, final ITestResult testResult) {
+    }
+
+    @Override
+    public void afterInvocation(final IInvokedMethod method, final ITestResult testResult) {
+        if (testResult.getStatus() != ITestResult.FAILURE) {
+            return;
+        }
+
+        final IRetryAnalyzer retryAnalyzer = testResult.getMethod().getRetryAnalyzer();
+        if (retryAnalyzer != null &&
+            retryAnalyzer instanceof FlakyRetryAnalyzer &&
+            !((FlakyRetryAnalyzer) retryAnalyzer).shouldRetry()) {
+            // Don't fail the build (flaky test), mark it as SKIPPED
+            testResult.setStatus(ITestResult.SKIP);
+            Reporter.setCurrentTestResult(testResult);
+        }
+    }
+}

diff --git a/util/src/test/java/org/killbill/billing/api/FlakyRetryAnalyzer.java b/util/src/test/java/org/killbill/billing/api/FlakyRetryAnalyzer.java
index 736e52d..9cacc19 100644
--- a/util/src/test/java/org/killbill/billing/api/FlakyRetryAnalyzer.java
+++ b/util/src/test/java/org/killbill/billing/api/FlakyRetryAnalyzer.java
@@ -32,13 +32,15 @@ public class FlakyRetryAnalyzer implements IRetryAnalyzer {
             return false;
         }
 
-        if (count < MAX_RETRIES) {
+        if (shouldRetry()) {
             count++;
             return true;
         } else {
-            // Don't fail the build (flaky test), mark it as SKIPPED
-            iTestResult.setStatus(ITestResult.SKIP);
             return false;
         }
     }
+
+    public boolean shouldRetry() {
+        return count < MAX_RETRIES;
+    }
 }

diff --git a/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuite.java b/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuite.java
index c7ba40e..6a60b18 100644
--- a/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuite.java
+++ b/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuite.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2016 Groupon, Inc
- * Copyright 2014-2016 The Billing Project, LLC
+ * Copyright 2014-2017 Groupon, Inc
+ * Copyright 2014-2017 The Billing Project, LLC
  *
  * The Billing Project 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
@@ -23,6 +23,7 @@ import java.util.UUID;
 
 import javax.inject.Inject;
 
+import org.killbill.billing.api.FlakyInvokedMethodListener;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.callcontext.MutableInternalCallContext;
 import org.killbill.billing.platform.api.KillbillConfigSource;
@@ -38,9 +39,11 @@ import org.slf4j.LoggerFactory;
 import org.testng.ITestResult;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Listeners;
 
 import com.google.common.collect.ImmutableMap;
 
+@Listeners(FlakyInvokedMethodListener.class)
 public class GuicyKillbillTestSuite {
 
     // Use the simple name here to save screen real estate