killbill-memoizeit

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeInitiated.java
index 6f04fd4..e720c55 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class AuthorizeInitiated extends PaymentLeavingStateCallback {
 
     public AuthorizeInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
index 5c2b1f9..e76c928 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
@@ -33,7 +33,7 @@ public class AuthorizeOperation extends PaymentOperation {
     public AuthorizeOperation(final PaymentAutomatonDAOHelper daoHelper,
                               final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                               final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureInitiated.java
index 78d5e4d..da7309b 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class CaptureInitiated extends PaymentLeavingStateCallback {
 
     public CaptureInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
index feaa39e..c3facf5 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
@@ -33,7 +33,7 @@ public class CaptureOperation extends PaymentOperation {
     public CaptureOperation(final PaymentAutomatonDAOHelper daoHelper,
                             final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                             final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackInitiated.java
index 72e9a3f..d4df130 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class ChargebackInitiated extends PaymentLeavingStateCallback {
 
     public ChargebackInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackOperation.java
index 6a272a6..14436a6 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/ChargebackOperation.java
@@ -19,8 +19,6 @@ package org.killbill.billing.payment.core.sm;
 
 import java.math.BigDecimal;
 
-import javax.annotation.Nullable;
-
 import org.killbill.automaton.OperationResult;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.TransactionType;
@@ -34,9 +32,6 @@ import org.killbill.commons.locker.GlobalLocker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.google.common.base.Predicate;
-import com.google.common.collect.Iterables;
-
 public class ChargebackOperation extends PaymentOperation {
 
     private final Logger logger = LoggerFactory.getLogger(ChargebackOperation.class);
@@ -44,7 +39,7 @@ public class ChargebackOperation extends PaymentOperation {
     public ChargebackOperation(final PaymentAutomatonDAOHelper daoHelper,
                                final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                                final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditInitiated.java
index 32ad101..befcbbd 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class CreditInitiated extends PaymentLeavingStateCallback {
 
     public CreditInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
index 1e3d550..1e6201e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
@@ -33,7 +33,7 @@ public class CreditOperation extends PaymentOperation {
     public CreditOperation(final PaymentAutomatonDAOHelper daoHelper,
                            final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                            final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonDAOHelper.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonDAOHelper.java
index 5e5cf0c..55fcb3f 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonDAOHelper.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonDAOHelper.java
@@ -134,14 +134,6 @@ public class PaymentAutomatonDAOHelper {
         paymentStateContext.setPaymentTransactionModelDao(paymentDao.getPaymentTransaction(paymentStateContext.getPaymentTransactionModelDao().getId(), internalCallContext));
     }
 
-    public UUID getDefaultPaymentMethodId() throws PaymentApiException {
-        final UUID paymentMethodId = paymentStateContext.getAccount().getPaymentMethodId();
-        if (paymentMethodId == null) {
-            throw new PaymentApiException(ErrorCode.PAYMENT_NO_DEFAULT_PAYMENT_METHOD, paymentStateContext.getAccount().getId());
-        }
-        return paymentMethodId;
-    }
-
     public PaymentPluginApi getPaymentProviderPlugin() throws PaymentApiException {
 
         final UUID paymentMethodId = paymentStateContext.getPaymentMethodId();

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
index 3b8068a..440ddde 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentAutomatonRunner.java
@@ -120,8 +120,9 @@ public class PaymentAutomatonRunner {
             Preconditions.checkState(currentStateName != null, "State name cannot be null for payment " + paymentId);
             Preconditions.checkState(paymentMethodId == null || effectivePaymentMethodId.equals(paymentMethodId), "Specified payment method id " + paymentMethodId + " doesn't match the one on the payment " + effectivePaymentMethodId);
         } else {
-            // If the payment method is not specified, retrieve the default one on the account
-            effectivePaymentMethodId = paymentMethodId != null ? paymentMethodId : daoHelper.getDefaultPaymentMethodId();
+            // If the payment method is not specified, retrieve the default one on the account; it could still be null, in which case
+            //
+            effectivePaymentMethodId = paymentMethodId != null ? paymentMethodId : account.getPaymentMethodId();
             currentStateName = paymentSMHelper.getInitStateNameForTransaction(transactionType);
         }
 

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentEnteringStateCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentEnteringStateCallback.java
index 430b132..a60afa2 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentEnteringStateCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentEnteringStateCallback.java
@@ -58,10 +58,10 @@ public abstract class PaymentEnteringStateCallback implements EnteringStateCallb
             final PaymentTransactionInfoPlugin paymentInfoPlugin = paymentStateContext.getPaymentInfoPlugin();
             final TransactionStatus paymentStatus = paymentPluginStatusToPaymentStatus(paymentInfoPlugin, operationResult);
             daoHelper.processPaymentInfoPlugin(paymentStatus, paymentInfoPlugin, newState.getName());
-        } else if (paymentStateContext.isApiPayment()) {
+        } else if (!paymentStateContext.isApiPayment()) {
             //
-            // If there is transaction to update (because payment transaction did not occur), we still want to send a bus event when the call originates from api
-            // to notify listeners that some transaction occurred for that specific account.
+            // If there is NO transaction to update (because payment transaction did not occur), then there is something wrong happening (maybe a missing defaultPaymentMethodId, ...)
+            // so, if the does NOT call originates from api then we still want to send a bus event so the system can react to it if needed.
             //
             final BusInternalEvent event = new DefaultPaymentErrorEvent(paymentStateContext.getAccount().getId(),
                                                                         null,

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentLeavingStateCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentLeavingStateCallback.java
index dcc8dc4..ad0419c 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentLeavingStateCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentLeavingStateCallback.java
@@ -20,6 +20,7 @@ package org.killbill.billing.payment.core.sm;
 import org.killbill.automaton.OperationException;
 import org.killbill.automaton.State;
 import org.killbill.automaton.State.LeavingStateCallback;
+import org.killbill.billing.ErrorCode;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -29,9 +30,11 @@ public abstract class PaymentLeavingStateCallback implements LeavingStateCallbac
     private final Logger logger = LoggerFactory.getLogger(PaymentLeavingStateCallback.class);
 
     protected final PaymentAutomatonDAOHelper daoHelper;
+    protected final PaymentStateContext paymentStateContext;
 
-    protected PaymentLeavingStateCallback(final PaymentAutomatonDAOHelper daoHelper) throws PaymentApiException {
+    protected PaymentLeavingStateCallback(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
         this.daoHelper = daoHelper;
+        this.paymentStateContext = paymentStateContext;
     }
 
     @Override
@@ -40,6 +43,12 @@ public abstract class PaymentLeavingStateCallback implements LeavingStateCallbac
 
         // Create or update the payment and transaction
         try {
+            // No paymentMethodId was passed through API and account does not have a default paymentMethodId
+            if (paymentStateContext.getPaymentMethodId() == null) {
+                throw new PaymentApiException(ErrorCode.PAYMENT_NO_DEFAULT_PAYMENT_METHOD, paymentStateContext.getAccount().getId());
+            }
+
+
             daoHelper.createNewPaymentTransaction();
         } catch (PaymentApiException e) {
             throw new OperationException(e);

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentOperation.java
index 4d4513e..be3da35 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentOperation.java
@@ -26,6 +26,7 @@ import org.killbill.automaton.Operation.OperationCallback;
 import org.killbill.automaton.OperationException;
 import org.killbill.automaton.OperationResult;
 import org.killbill.billing.ErrorCode;
+import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.TransactionStatus;
 import org.killbill.billing.payment.api.TransactionType;
@@ -48,24 +49,32 @@ import com.google.common.collect.Iterables;
 // Encapsulates the payment specific logic
 public abstract class PaymentOperation extends OperationCallbackBase implements OperationCallback {
 
-    protected final PaymentPluginApi plugin;
     protected final PaymentAutomatonDAOHelper daoHelper;
+    protected PaymentPluginApi plugin;
 
-    protected PaymentOperation(final PaymentAutomatonDAOHelper daoHelper, final GlobalLocker locker,
+    protected PaymentOperation(final GlobalLocker locker,
+                               final PaymentAutomatonDAOHelper daoHelper,
                                final PluginDispatcher<OperationResult> paymentPluginDispatcher,
-                               final PaymentStateContext paymentStateContext) throws PaymentApiException {
+                               final PaymentStateContext paymentStateContext) {
         super(locker, paymentPluginDispatcher, paymentStateContext);
         this.daoHelper = daoHelper;
-        this.plugin = daoHelper.getPaymentProviderPlugin();
     }
 
 
     @Override
     public OperationResult doOperationCallback() throws OperationException {
-        if (paymentStateContext.shouldLockAccountAndDispatch()) {
-            return doOperationCallbackWithDispatchAndAccountLock();
-        } else {
-            return doSimpleOperationCallback();
+
+        try {
+
+            this.plugin = daoHelper.getPaymentProviderPlugin();
+
+            if (paymentStateContext.shouldLockAccountAndDispatch()) {
+                return doOperationCallbackWithDispatchAndAccountLock();
+            } else {
+                return doSimpleOperationCallback();
+            }
+        } catch (PaymentApiException e) {
+            throw new OperationException(e, OperationResult.EXCEPTION);
         }
     }
 

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseInitiated.java
index edb69e5..c5e97e1 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class PurchaseInitiated extends PaymentLeavingStateCallback {
 
     public PurchaseInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseOperation.java
index 2c835c9..39ef3c6 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PurchaseOperation.java
@@ -33,7 +33,7 @@ public class PurchaseOperation extends PaymentOperation {
     public PurchaseOperation(final PaymentAutomatonDAOHelper daoHelper,
                              final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                              final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundInitiated.java
index f9000ec..095313c 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class RefundInitiated extends PaymentLeavingStateCallback {
 
     public RefundInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
index d977299..b78700e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
@@ -33,7 +33,7 @@ public class RefundOperation extends PaymentOperation {
     public RefundOperation(final PaymentAutomatonDAOHelper daoHelper,
                            final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                            final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidInitiated.java
index 89e8b03..b5bbb65 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidInitiated.java
@@ -22,6 +22,6 @@ import org.killbill.billing.payment.api.PaymentApiException;
 public class VoidInitiated extends PaymentLeavingStateCallback {
 
     public VoidInitiated(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper);
+        super(daoHelper, paymentStateContext);
     }
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
index c4871f5..5427c27 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
@@ -33,7 +33,7 @@ public class VoidOperation extends PaymentOperation {
     public VoidOperation(final PaymentAutomatonDAOHelper daoHelper,
                          final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher,
                          final PaymentStateContext paymentStateContext) throws PaymentApiException {
-        super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+        super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
     }
 
     @Override

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentAutomatonDAOHelper.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentAutomatonDAOHelper.java
index 140780d..a5f96fa 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentAutomatonDAOHelper.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentAutomatonDAOHelper.java
@@ -97,17 +97,6 @@ public class TestPaymentAutomatonDAOHelper extends PaymentTestSuiteWithEmbeddedD
     }
 
     @Test(groups = "slow")
-    public void testNoDefaultPaymentMethod() throws Exception {
-        final PaymentAutomatonDAOHelper daoHelper = createDAOHelper(UUID.randomUUID(), paymentExternalKey, paymentTransactionExternalKey, amount, currency);
-        try {
-            daoHelper.getDefaultPaymentMethodId();
-            Assert.fail();
-        } catch (final PaymentApiException e) {
-            Assert.assertEquals(e.getCode(), ErrorCode.PAYMENT_NO_DEFAULT_PAYMENT_METHOD.getCode());
-        }
-    }
-
-    @Test(groups = "slow")
     public void testNoPaymentMethod() throws Exception {
         final PaymentAutomatonDAOHelper daoHelper = createDAOHelper(UUID.randomUUID(), paymentExternalKey, paymentTransactionExternalKey, amount, currency);
         try {

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
index 17e50c9..172a623 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
@@ -136,15 +136,15 @@ public class TestPaymentLeavingStateCallback extends PaymentTestSuiteWithEmbedde
         }
 
         final PaymentAutomatonDAOHelper daoHelper = new PaymentAutomatonDAOHelper(paymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext, eventBus, paymentSMHelper);
-        callback = new PaymentLeavingStateTestCallback(daoHelper);
+        callback = new PaymentLeavingStateTestCallback(daoHelper, paymentStateContext);
 
         Mockito.when(state.getName()).thenReturn("NEW_STATE");
     }
 
     private static final class PaymentLeavingStateTestCallback extends PaymentLeavingStateCallback {
 
-        private PaymentLeavingStateTestCallback(final PaymentAutomatonDAOHelper daoHelper) throws PaymentApiException {
-            super(daoHelper);
+        private PaymentLeavingStateTestCallback(final PaymentAutomatonDAOHelper daoHelper, final PaymentStateContext paymentStateContext) throws PaymentApiException {
+            super(daoHelper, paymentStateContext);
         }
     }
 }

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentOperation.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentOperation.java
index a21aefd..f8c14a5 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentOperation.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentOperation.java
@@ -134,7 +134,7 @@ public class TestPaymentOperation extends PaymentTestSuiteNoDB {
         public PaymentOperationTest(@Nullable final PaymentPluginStatus paymentPluginStatus,
                                     final PaymentAutomatonDAOHelper daoHelper, final GlobalLocker locker,
                                     final PluginDispatcher<OperationResult> paymentPluginDispatcher, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-            super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+            super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
             this.paymentInfoPlugin = (paymentPluginStatus == null ? null : getPaymentInfoPlugin(paymentPluginStatus));
         }
 

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
index 4362013..f6aae85 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
@@ -291,7 +291,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
     private static final class PluginOperationTest extends PaymentOperation {
 
         protected PluginOperationTest(final PaymentAutomatonDAOHelper daoHelper, final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher, final PaymentStateContext paymentStateContext) throws PaymentApiException {
-            super(daoHelper, locker, paymentPluginDispatcher, paymentStateContext);
+            super(locker, daoHelper, paymentPluginDispatcher, paymentStateContext);
         }
 
         @Override