killbill-memoizeit

Introduce PaymentStateMachineHelper and RetryStateMachineHelper

7/10/2014 5:19:34 PM

Changes

Details

diff --git a/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java b/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
index 0c621ac..746204c 100644
--- a/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
+++ b/payment/src/main/java/org/killbill/billing/payment/control/InvoicePaymentControlPluginApi.java
@@ -183,9 +183,7 @@ public final class InvoicePaymentControlPluginApi implements PaymentControlPlugi
                     throw new IllegalStateException("Unexpected transactionType " + transactionType);
             }
         } catch (InvoiceApiException e) {
-            // STEPH need to add some state machine logic in the plugin itself to handle those cases
-            logger.error("Failed to complete call: ", e);
-            //throw new PaymentControlApiException(e);
+            logger.error("InvoicePaymentControlPluginApi onSuccessCall failed for attemptId = " + paymentControlContext.getAttemptPaymentId() + ", transactionType  = " + transactionType , e);
         }
     }
 
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/DirectPaymentProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/DirectPaymentProcessor.java
index 9f05a04..cae84dd 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/DirectPaymentProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/DirectPaymentProcessor.java
@@ -50,6 +50,7 @@ import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionStatus;
 import org.killbill.billing.payment.api.TransactionType;
 import org.killbill.billing.payment.core.sm.DirectPaymentAutomatonRunner;
+import org.killbill.billing.payment.core.sm.PaymentStateMachineHelper;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.dao.PaymentModelDao;
 import org.killbill.billing.payment.dao.PaymentTransactionModelDao;
@@ -87,6 +88,7 @@ public class DirectPaymentProcessor extends ProcessorBase {
 
     private final DirectPaymentAutomatonRunner directPaymentAutomatonRunner;
     private final InternalCallContextFactory internalCallContextFactory;
+    private final PaymentStateMachineHelper paymentSMHelper;
 
     private static final Logger log = LoggerFactory.getLogger(DirectPaymentProcessor.class);
 
@@ -102,8 +104,10 @@ public class DirectPaymentProcessor extends ProcessorBase {
                                   final GlobalLocker locker,
                                   @Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor,
                                   final DirectPaymentAutomatonRunner directPaymentAutomatonRunner,
+                                  final PaymentStateMachineHelper paymentSMHelper,
                                   final Clock clock) {
         super(pluginRegistry, accountUserApi, eventBus, paymentDao, nonEntityDao, tagUserApi, locker, executor, invoiceApi, clock);
+        this.paymentSMHelper = paymentSMHelper;
         this.internalCallContextFactory = internalCallContextFactory;
         this.directPaymentAutomatonRunner = directPaymentAutomatonRunner;
     }
@@ -175,13 +179,10 @@ public class DirectPaymentProcessor extends ProcessorBase {
 
         final TransactionStatus newStatus = isSuccess ? TransactionStatus.SUCCESS : TransactionStatus.PAYMENT_FAILURE;
         // STEPH This works if the pending transaction we are trying to update matches is the one that gave the state to the payment. Also can we have multiple PENDING for a given payment?
-        final State currentPaymentState = directPaymentAutomatonRunner.fetchNextState(paymentModelDao.getStateName(), isSuccess);
-        // STEPH : should we insert a new transaction row to keep the PENDING one?
-
-        // STEPH hack; need proper automaton API to understand what is a successful terminal state.
-        final String lastSuccessPaymentStateStrOrNull = currentPaymentState.getName().endsWith("SUCCESS") ? currentPaymentState.getName() : null;
-
-        paymentDao.updateDirectPaymentAndTransactionOnCompletion(transactionModelDao.getPaymentId(), currentPaymentState.getName(), lastSuccessPaymentStateStrOrNull, transactionModelDao.getId(), newStatus,
+        final State currentPaymentState;
+        final String stateName = paymentModelDao.getStateName();
+        final String lastSuccessPaymentStateStrOrNull = paymentSMHelper.isSuccessState(stateName) ? stateName : null;
+        paymentDao.updateDirectPaymentAndTransactionOnCompletion(transactionModelDao.getPaymentId(), stateName, lastSuccessPaymentStateStrOrNull, transactionModelDao.getId(), newStatus,
                                                                  transactionModelDao.getProcessedAmount(), transactionModelDao.getProcessedCurrency(),
                                                                  transactionModelDao.getGatewayErrorCode(), transactionModelDao.getGatewayErrorMsg(), internalCallContext);
     }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/Janitor.java b/payment/src/main/java/org/killbill/billing/payment/core/Janitor.java
index 904ce77..2d7cb57 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/Janitor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/Janitor.java
@@ -37,6 +37,7 @@ import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.TransactionStatus;
 import org.killbill.billing.payment.core.sm.PluginControlledDirectPaymentAutomatonRunner;
+import org.killbill.billing.payment.core.sm.RetryStateMachineHelper;
 import org.killbill.billing.payment.core.sm.RetryableDirectPaymentStateContext;
 import org.killbill.billing.payment.dao.PaymentAttemptModelDao;
 import org.killbill.billing.payment.dao.PaymentDao;
@@ -72,6 +73,7 @@ public class Janitor {
     private final InternalCallContextFactory internalCallContextFactory;
     private final NonEntityDao nonEntityDao;
     private final PluginControlledDirectPaymentAutomatonRunner pluginControlledDirectPaymentAutomatonRunner;
+    private final RetryStateMachineHelper retrySMHelper;
 
     private volatile boolean isStopped;
     private CountDownLatch shutdownLatch;
@@ -84,7 +86,8 @@ public class Janitor {
                    final NonEntityDao nonEntityDao,
                    final InternalCallContextFactory internalCallContextFactory,
                    final PluginControlledDirectPaymentAutomatonRunner pluginControlledDirectPaymentAutomatonRunner,
-                   @Named(PaymentModule.JANITOR_EXECUTOR_NAMED) final ScheduledExecutorService janitorExecutor) {
+                   @Named(PaymentModule.JANITOR_EXECUTOR_NAMED) final ScheduledExecutorService janitorExecutor,
+                   final RetryStateMachineHelper retrySMHelper) {
         this.accountInternalApi = accountInternalApi;
         this.paymentDao = paymentDao;
         this.clock = clock;
@@ -93,6 +96,7 @@ public class Janitor {
         this.nonEntityDao = nonEntityDao;
         this.internalCallContextFactory = internalCallContextFactory;
         this.pluginControlledDirectPaymentAutomatonRunner = pluginControlledDirectPaymentAutomatonRunner;
+        this.retrySMHelper = retrySMHelper;
     }
 
     public void start() {
@@ -176,8 +180,7 @@ public class Janitor {
                 return;
             }
 
-            // STEPH state string hack
-            final List<PaymentAttemptModelDao> incompleteAttempts = paymentDao.getPaymentAttemptsByState("INIT", getCreatedDateBefore(), fakeCallContext);
+            final List<PaymentAttemptModelDao> incompleteAttempts = paymentDao.getPaymentAttemptsByState(retrySMHelper.getInitialState().getName(), getCreatedDateBefore(), fakeCallContext);
             log.info("AttemptCompletionTask start run : found " + incompleteAttempts.size() + " incomplete attempts");
 
             for (PaymentAttemptModelDao cur : incompleteAttempts) {
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/PluginControlledPaymentProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/PluginControlledPaymentProcessor.java
index c733015..2ef76f6 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/PluginControlledPaymentProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/PluginControlledPaymentProcessor.java
@@ -23,6 +23,7 @@ import java.util.concurrent.ExecutorService;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 
+import org.killbill.automaton.MissingEntryException;
 import org.killbill.automaton.State;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.account.api.Account;
@@ -37,6 +38,7 @@ import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionType;
 import org.killbill.billing.payment.core.sm.PluginControlledDirectPaymentAutomatonRunner;
+import org.killbill.billing.payment.core.sm.RetryStateMachineHelper;
 import org.killbill.billing.payment.dao.PaymentAttemptModelDao;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.dao.PaymentModelDao;
@@ -58,6 +60,7 @@ import static org.killbill.billing.payment.glue.PaymentModule.PLUGIN_EXECUTOR_NA
 public class PluginControlledPaymentProcessor extends ProcessorBase {
 
     private final PluginControlledDirectPaymentAutomatonRunner pluginControlledDirectPaymentAutomatonRunner;
+    private final RetryStateMachineHelper retrySMHelper;
 
     @Inject
     public PluginControlledPaymentProcessor(final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry,
@@ -70,9 +73,10 @@ public class PluginControlledPaymentProcessor extends ProcessorBase {
                                             final GlobalLocker locker,
                                             @Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor,
                                             final PluginControlledDirectPaymentAutomatonRunner pluginControlledDirectPaymentAutomatonRunner,
+                                            final RetryStateMachineHelper retrySMHelper,
                                             final Clock clock) {
         super(pluginRegistry, accountInternalApi, eventBus, paymentDao, nonEntityDao, tagUserApi, locker, executor, invoiceApi, clock);
-
+        this.retrySMHelper = retrySMHelper;
         this.pluginControlledDirectPaymentAutomatonRunner = pluginControlledDirectPaymentAutomatonRunner;
     }
 
@@ -204,7 +208,7 @@ public class PluginControlledPaymentProcessor extends ProcessorBase {
             final UUID tenantId = nonEntityDao.retrieveIdFromObject(internalCallContext.getTenantRecordId(), ObjectType.TENANT);
             final CallContext callContext = internalCallContext.toCallContext(tenantId);
 
-            final State state = pluginControlledDirectPaymentAutomatonRunner.fetchState(attempt.getStateName());
+            final State state = retrySMHelper.getState(attempt.getStateName());
             pluginControlledDirectPaymentAutomatonRunner.run(state,
                                                              false,
                                                              attempt.getTransactionType(),
@@ -226,6 +230,8 @@ public class PluginControlledPaymentProcessor extends ProcessorBase {
             log.warn("Failed to retry attempt " + attemptId + " for plugin " + pluginName, e);
         } catch (PluginPropertySerializerException e) {
             log.warn("Failed to retry attempt " + attemptId + " for plugin " + pluginName, e);
+        } catch (MissingEntryException e) {
+            log.warn("Failed to retry attempt " + attemptId + " for plugin " + pluginName, e);
         }
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonDAOHelper.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonDAOHelper.java
index 4be2d4c..b0347ea 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonDAOHelper.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonDAOHelper.java
@@ -44,6 +44,7 @@ public class DirectPaymentAutomatonDAOHelper {
     protected final DirectPaymentStateContext directPaymentStateContext;
     protected final DateTime utcNow;
     protected final InternalCallContext internalCallContext;
+    protected final PaymentStateMachineHelper paymentSMHelper;
 
     protected final PaymentDao paymentDao;
 
@@ -53,12 +54,14 @@ public class DirectPaymentAutomatonDAOHelper {
     public DirectPaymentAutomatonDAOHelper(final DirectPaymentStateContext directPaymentStateContext,
                                            final DateTime utcNow, final PaymentDao paymentDao,
                                            final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry,
-                                           final InternalCallContext internalCallContext) throws PaymentApiException {
+                                           final InternalCallContext internalCallContext,
+                                           final PaymentStateMachineHelper paymentSMHelper) throws PaymentApiException {
         this.directPaymentStateContext = directPaymentStateContext;
         this.utcNow = utcNow;
         this.paymentDao = paymentDao;
         this.pluginRegistry = pluginRegistry;
         this.internalCallContext = internalCallContext;
+        this.paymentSMHelper = paymentSMHelper;
     }
 
     public void createNewDirectPaymentTransaction() throws PaymentApiException {
@@ -97,9 +100,7 @@ public class DirectPaymentAutomatonDAOHelper {
         final String gatewayErrorCode = paymentInfoPlugin == null ? null : paymentInfoPlugin.getGatewayErrorCode();
         final String gatewayErrorMsg = paymentInfoPlugin == null ? null : paymentInfoPlugin.getGatewayError();
 
-        // STEPH hack
-        final String lastSuccessPaymentState = currentPaymentStateName.endsWith("SUCCESS") ? currentPaymentStateName : null;
-
+        final String lastSuccessPaymentState = paymentSMHelper.isSuccessState(currentPaymentStateName) ? currentPaymentStateName : null;
         paymentDao.updateDirectPaymentAndTransactionOnCompletion(directPaymentStateContext.getDirectPaymentId(),
                                                                  currentPaymentStateName,
                                                                  lastSuccessPaymentState,
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
index b13ada1..46b5552 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
@@ -45,8 +45,8 @@ import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionType;
-import org.killbill.billing.payment.dao.PaymentModelDao;
 import org.killbill.billing.payment.dao.PaymentDao;
+import org.killbill.billing.payment.dao.PaymentModelDao;
 import org.killbill.billing.payment.dispatcher.PluginDispatcher;
 import org.killbill.billing.payment.glue.PaymentModule;
 import org.killbill.billing.payment.plugin.api.PaymentPluginApi;
@@ -66,7 +66,7 @@ import static org.killbill.billing.payment.glue.PaymentModule.PLUGIN_EXECUTOR_NA
 
 public class DirectPaymentAutomatonRunner {
 
-    protected final StateMachineConfig stateMachineConfig;
+    protected final PaymentStateMachineHelper paymentSMHelper;
     protected final PaymentDao paymentDao;
     protected final GlobalLocker locker;
     protected final PluginDispatcher<OperationResult> paymentPluginDispatcher;
@@ -80,8 +80,9 @@ public class DirectPaymentAutomatonRunner {
                                         final GlobalLocker locker,
                                         final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry,
                                         final Clock clock,
-                                        @Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor) {
-        this.stateMachineConfig = stateMachineConfig;
+                                        @Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor,
+                                        final PaymentStateMachineHelper paymentSMHelper) {
+        this.paymentSMHelper = paymentSMHelper;
         this.paymentDao = paymentDao;
         this.locker = locker;
         this.pluginRegistry = pluginRegistry;
@@ -102,7 +103,7 @@ public class DirectPaymentAutomatonRunner {
 
         final DirectPaymentStateContext directPaymentStateContext = new DirectPaymentStateContext(directPaymentId, attemptId, directPaymentExternalKey, directPaymentTransactionExternalKey, transactionType,
                                                                                                   account, paymentMethodId, amount, currency, shouldLockAccount, properties, internalCallContext, callContext);
-        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, utcNow, paymentDao, pluginRegistry, internalCallContext);
+        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, utcNow, paymentDao, pluginRegistry, internalCallContext, paymentSMHelper);
 
         final UUID effectivePaymentMethodId;
         final String currentStateMachineName;
@@ -110,18 +111,15 @@ public class DirectPaymentAutomatonRunner {
         if (directPaymentId != null) {
             final PaymentModelDao paymentModelDao = daoHelper.getDirectPayment();
             effectivePaymentMethodId = paymentModelDao.getPaymentMethodId();
-            currentStateName = paymentModelDao.getLastSuccessStateName() != null ? paymentModelDao.getLastSuccessStateName() : getInitState(transactionType);
-            currentStateMachineName = getStateMachineName(currentStateName);
+            currentStateName = paymentModelDao.getLastSuccessStateName() != null ? paymentModelDao.getLastSuccessStateName() : paymentSMHelper.getInitStateNameForTransaction(transactionType);
 
             // Check for illegal states (should never happen)
-            Preconditions.checkState(currentStateMachineName != null, "State machine name cannot be null for direct payment " + directPaymentId);
             Preconditions.checkState(currentStateName != null, "State name cannot be null for direct payment " + directPaymentId);
             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();
-            currentStateName = getInitState(transactionType);
-            currentStateMachineName = getStateMachineName(currentStateName);
+            currentStateName = paymentSMHelper.getInitStateNameForTransaction(transactionType);
         }
 
         directPaymentStateContext.setPaymentMethodId(effectivePaymentMethodId);
@@ -133,50 +131,36 @@ public class DirectPaymentAutomatonRunner {
         final EnteringStateCallback enteringStateCallback;
         switch (transactionType) {
             case PURCHASE:
-                operationStateMachineName = "PURCHASE";
-                operationName = "OP_PURCHASE";
                 operationCallback = new PurchaseOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new PurchaseInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new PurchaseCompleted(daoHelper, directPaymentStateContext);
                 break;
             case AUTHORIZE:
-                operationStateMachineName = "AUTHORIZE";
-                operationName = "OP_AUTHORIZE";
                 operationCallback = new AuthorizeOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new AuthorizeInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new AuthorizeCompleted(daoHelper, directPaymentStateContext);
                 break;
             case CAPTURE:
-                operationStateMachineName = "CAPTURE";
-                operationName = "OP_CAPTURE";
                 operationCallback = new CaptureOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new CaptureInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new CaptureCompleted(daoHelper, directPaymentStateContext);
                 break;
             case VOID:
-                operationStateMachineName = "VOID";
-                operationName = "OP_VOID";
                 operationCallback = new VoidOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new VoidInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new VoidCompleted(daoHelper, directPaymentStateContext);
                 break;
             case REFUND:
-                operationStateMachineName = "REFUND";
-                operationName = "OP_REFUND";
                 operationCallback = new RefundOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new RefundInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new RefundCompleted(daoHelper, directPaymentStateContext);
                 break;
             case CREDIT:
-                operationStateMachineName = "CREDIT";
-                operationName = "OP_CREDIT";
                 operationCallback = new CreditOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new CreditInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new CreditCompleted(daoHelper, directPaymentStateContext);
                 break;
             case CHARGEBACK:
-                operationStateMachineName = "CHARGEBACK";
-                operationName = "OP_CHARGEBACK";
                 operationCallback = new ChargebackOperation(daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
                 leavingStateCallback = new ChargebackInitiated(daoHelper, directPaymentStateContext);
                 enteringStateCallback = new ChargebackCompleted(daoHelper, directPaymentStateContext);
@@ -185,13 +169,13 @@ public class DirectPaymentAutomatonRunner {
                 throw new IllegalStateException("Unsupported transaction type " + transactionType);
         }
 
-        runStateMachineOperation(currentStateMachineName, currentStateName, operationStateMachineName, operationName, leavingStateCallback, operationCallback, enteringStateCallback);
+        runStateMachineOperation(currentStateName, transactionType, leavingStateCallback, operationCallback, enteringStateCallback);
 
         return directPaymentStateContext.getDirectPaymentId();
     }
 
-    public final State fetchNextState(final String prevStateName, final boolean isSuccess) {
-        final StateMachine stateMachine = getStateMachine(prevStateName);
+    public final State fetchNextState(final String prevStateName, final boolean isSuccess) throws MissingEntryException {
+        final StateMachine stateMachine = paymentSMHelper.getStateMachineForStateName(prevStateName);
         final Transition transition = Iterables.tryFind(ImmutableList.copyOf(stateMachine.getTransitions()), new Predicate<Transition>() {
             @Override
             public boolean apply(final Transition input) {
@@ -203,48 +187,12 @@ public class DirectPaymentAutomatonRunner {
         return transition != null ? transition.getFinalState() : null;
     }
 
-    // Hack for now
-    protected String getStateMachineName(final String currentStateName) {
-        final StateMachine stateMachine = getStateMachine(currentStateName);
-        if (stateMachine == null) {
-            return null;
-        }
-        return stateMachine.getName();
-    }
-
-    private String getInitState(final TransactionType transactionType) {
-        switch (transactionType) {
-            case AUTHORIZE:
-                return "AUTH_INIT";
-            case CREDIT:
-                return "CREDIT_INIT";
-            case PURCHASE:
-                return "PURCHASE_INIT";
-            default:
-                throw new IllegalStateException("Unsupported transaction type " + transactionType + " for null direct payment id");
-        }
-    }
-
-    private StateMachine getStateMachine(final String currentStateName) {
-        for (final StateMachine stateMachine : stateMachineConfig.getStateMachines()) {
-            for (final State state : stateMachine.getStates()) {
-                if (state.getName().equals(currentStateName)) {
-                    return stateMachine;
-                }
-            }
-        }
-        return null;
-    }
-
-    protected void runStateMachineOperation(final String initialStateMachineName, final String initialStateName,
-                                            final String operationStateMachineName, final String operationName,
+    protected void runStateMachineOperation(final String initialStateName, final TransactionType transactionType,
                                             final LeavingStateCallback leavingStateCallback, final OperationCallback operationCallback, final EnteringStateCallback enteringStateCallback) throws PaymentApiException {
         try {
-            final StateMachine initialStateMachine = stateMachineConfig.getStateMachine(initialStateMachineName);
+            final StateMachine initialStateMachine = paymentSMHelper.getStateMachineForStateName(initialStateName);
             final State initialState = initialStateMachine.getState(initialStateName);
-
-            final StateMachine operationStateMachine = stateMachineConfig.getStateMachine(operationStateMachineName);
-            final Operation operation = operationStateMachine.getOperation(operationName);
+            final Operation operation = paymentSMHelper.getOperationForTransaction(transactionType);
 
             initialState.runOperation(operation, operationCallback, enteringStateCallback, leavingStateCallback);
         } catch (final MissingEntryException e) {
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateMachineHelper.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateMachineHelper.java
new file mode 100644
index 0000000..665e1ca
--- /dev/null
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateMachineHelper.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright 2014 Groupon, Inc
+ * Copyright 2014 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.payment.core.sm;
+
+import javax.inject.Inject;
+
+import org.killbill.automaton.MissingEntryException;
+import org.killbill.automaton.Operation;
+import org.killbill.automaton.OperationResult;
+import org.killbill.automaton.State;
+import org.killbill.automaton.StateMachine;
+import org.killbill.automaton.StateMachineConfig;
+import org.killbill.automaton.Transition;
+import org.killbill.billing.payment.api.TransactionType;
+import org.killbill.billing.payment.glue.PaymentModule;
+
+import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
+
+public class PaymentStateMachineHelper {
+
+    private static final String AUTHORIZE_STATE_MACHINE_NAME = "AUTHORIZE";
+    private static final String CAPTURE_STATE_MACHINE_NAME = "CAPTURE";
+    private static final String PURCHASE_STATE_MACHINE_NAME = "PURCHASE";
+    private static final String REFUND_STATE_MACHINE_NAME = "REFUND";
+    private static final String CREDIT_STATE_MACHINE_NAME = "CREDIT";
+    private static final String VOID_STATE_MACHINE_NAME = "VOID";
+    private static final String CHARGEBACK_STATE_MACHINE_NAME = "CHARGEBACK";
+
+    private static final String AUTHORIZE_INIT_STATE_NAME = "AUTH_INIT";
+    private static final String CAPTURE_INIT_STATE_NAME = "CAPTURE_INIT";
+    private static final String PURCHASE_INIT_STATE_NAME = "PURCHASE_INIT";
+    private static final String REFUND_INIT_STATE_NAME = "REFUND_INIT";
+    private static final String CREDIT_INIT_STATE_NAME = "CREDIT_INIT";
+    private static final String VOID_INIT_STATE_NAME = "VOID_INIT";
+    private static final String CHARGEBACK_INIT_STATE_NAME = "CHARGEBACK_INIT";
+
+    private final StateMachineConfig stateMachineConfig;
+
+    @Inject
+    public PaymentStateMachineHelper(@javax.inject.Named(PaymentModule.STATE_MACHINE_PAYMENT) final StateMachineConfig stateMachineConfig) {
+        this.stateMachineConfig = stateMachineConfig;
+    }
+
+    public State getState(final String stateName) throws MissingEntryException {
+        final StateMachine stateMachine  = stateMachineConfig.getStateMachineForState(stateName);
+        return stateMachine.getState(stateName);
+    }
+
+    public String getInitStateNameForTransaction(final TransactionType transactionType) {
+        switch (transactionType) {
+            case AUTHORIZE:
+                return AUTHORIZE_INIT_STATE_NAME;
+            case CAPTURE:
+                return CAPTURE_INIT_STATE_NAME;
+            case PURCHASE:
+                return PURCHASE_INIT_STATE_NAME;
+            case REFUND:
+                return REFUND_INIT_STATE_NAME;
+            case CREDIT:
+                return CREDIT_INIT_STATE_NAME;
+            case VOID:
+                return VOID_INIT_STATE_NAME;
+            case CHARGEBACK:
+                return CHARGEBACK_INIT_STATE_NAME;
+            default:
+                throw new IllegalStateException("Unsupported transaction type " + transactionType + " for null direct payment id");
+        }
+    }
+
+    public StateMachine getStateMachineForStateName(final String stateName) throws MissingEntryException {
+        return stateMachineConfig.getStateMachineForState(stateName);
+    }
+
+    public Operation getOperationForTransaction(final TransactionType transactionType) throws MissingEntryException {
+        final StateMachine stateMachine = getStateMachineForTransaction(transactionType);
+        // Only one operation defined, this is the current PaymentStates.xml model
+        return stateMachine.getOperations()[0];
+    }
+
+    public StateMachine getStateMachineForTransaction(final TransactionType transactionType) throws MissingEntryException {
+        switch (transactionType) {
+            case AUTHORIZE:
+                return stateMachineConfig.getStateMachine(AUTHORIZE_STATE_MACHINE_NAME);
+            case CAPTURE:
+                return stateMachineConfig.getStateMachine(CAPTURE_STATE_MACHINE_NAME);
+            case PURCHASE:
+                return stateMachineConfig.getStateMachine(PURCHASE_STATE_MACHINE_NAME);
+            case REFUND:
+                return stateMachineConfig.getStateMachine(REFUND_STATE_MACHINE_NAME);
+            case CREDIT:
+                return stateMachineConfig.getStateMachine(CREDIT_STATE_MACHINE_NAME);
+            case VOID:
+                return stateMachineConfig.getStateMachine(VOID_STATE_MACHINE_NAME);
+            case CHARGEBACK:
+                return stateMachineConfig.getStateMachine(CHARGEBACK_STATE_MACHINE_NAME);
+            default:
+                throw new IllegalStateException("Unsupported transaction type " + transactionType + " for null direct payment id");
+        }
+    }
+
+    // A better way would be to change the xml to add attributes to the state (e.g isTerminal, isSuccess, isInit,...)
+    public boolean isSuccessState(final String stateName) {
+        return stateName.endsWith("SUCCESS");
+    }
+
+    public final State fetchNextState(final String prevStateName, final boolean isSuccess) throws MissingEntryException {
+        final StateMachine stateMachine = getStateMachineForStateName(prevStateName);
+        final Transition transition = Iterables.tryFind(ImmutableList.copyOf(stateMachine.getTransitions()), new Predicate<Transition>() {
+            @Override
+            public boolean apply(final Transition input) {
+                // This works because there is only one operation defined for a given state machine, which is our model for PaymentStates.xml
+                return input.getInitialState().getName().equals(prevStateName) &&
+                       input.getOperationResult().equals(isSuccess ? OperationResult.SUCCESS : OperationResult.FAILURE);
+            }
+        }).orNull();
+        return transition != null ? transition.getFinalState() : null;
+    }
+}
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PluginControlledDirectPaymentAutomatonRunner.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PluginControlledDirectPaymentAutomatonRunner.java
index 31a6376..cc265c4 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PluginControlledDirectPaymentAutomatonRunner.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PluginControlledDirectPaymentAutomatonRunner.java
@@ -25,13 +25,11 @@ import javax.inject.Inject;
 import javax.inject.Named;
 
 import org.killbill.automaton.MissingEntryException;
-import org.killbill.automaton.Operation;
 import org.killbill.automaton.Operation.OperationCallback;
 import org.killbill.automaton.OperationException;
 import org.killbill.automaton.State;
 import org.killbill.automaton.State.EnteringStateCallback;
 import org.killbill.automaton.State.LeavingStateCallback;
-import org.killbill.automaton.StateMachine;
 import org.killbill.automaton.StateMachineConfig;
 import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.Account;
@@ -62,37 +60,30 @@ import static org.killbill.billing.payment.glue.PaymentModule.RETRYABLE_NAMED;
 
 public class PluginControlledDirectPaymentAutomatonRunner extends DirectPaymentAutomatonRunner {
 
-    private final StateMachine retryStateMachine;
-
     private final DirectPaymentProcessor directPaymentProcessor;
-    private final State initialState;
-    private final State retriedState;
-    private final Operation retryOperation;
     private final RetryServiceScheduler retryServiceScheduler;
 
+    private final RetryStateMachineHelper retrySMHelper;
+
     protected final OSGIServiceRegistration<PaymentControlPluginApi> paymentControlPluginRegistry;
 
     @Inject
-    public PluginControlledDirectPaymentAutomatonRunner(@Named(PaymentModule.STATE_MACHINE_PAYMENT) final StateMachineConfig stateMachineConfig, @Named(PaymentModule.STATE_MACHINE_RETRY) final StateMachineConfig retryStateMachine, final PaymentDao paymentDao, final GlobalLocker locker, final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry,
-                                                        final OSGIServiceRegistration<PaymentControlPluginApi> retryPluginRegistry, final Clock clock, final TagInternalApi tagApi, final DirectPaymentProcessor directPaymentProcessor, @Named(RETRYABLE_NAMED) final RetryServiceScheduler retryServiceScheduler, final PaymentConfig paymentConfig,
-                                                        @com.google.inject.name.Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor) {
-        super(stateMachineConfig, paymentConfig, paymentDao, locker, pluginRegistry, clock, executor);
+    public PluginControlledDirectPaymentAutomatonRunner(@Named(PaymentModule.STATE_MACHINE_PAYMENT) final StateMachineConfig stateMachineConfig, final PaymentDao paymentDao, final GlobalLocker locker, final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry,
+                                                        final OSGIServiceRegistration<PaymentControlPluginApi> retryPluginRegistry, final Clock clock, final DirectPaymentProcessor directPaymentProcessor, @Named(RETRYABLE_NAMED) final RetryServiceScheduler retryServiceScheduler,
+                                                        final PaymentConfig paymentConfig, @com.google.inject.name.Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor, PaymentStateMachineHelper paymentSMHelper, RetryStateMachineHelper retrySMHelper) {
+        super(stateMachineConfig, paymentConfig, paymentDao, locker, pluginRegistry, clock, executor, paymentSMHelper);
         this.directPaymentProcessor = directPaymentProcessor;
         this.paymentControlPluginRegistry = retryPluginRegistry;
         this.retryServiceScheduler = retryServiceScheduler;
-        this.retryStateMachine = retryStateMachine.getStateMachines()[0];
-        this.initialState = fetchState("INIT");
-        this.retriedState = fetchState("RETRIED");
-        this.retryOperation = fetchRetryOperation();
+        this.retrySMHelper = retrySMHelper;
     }
 
-
     public DirectPayment run(final boolean isApiPayment, final TransactionType transactionType, final Account account, @Nullable final UUID paymentMethodId,
                              @Nullable final UUID directPaymentId, @Nullable final String directPaymentExternalKey, final String directPaymentTransactionExternalKey,
                              @Nullable final BigDecimal amount, @Nullable final Currency currency,
                              final Iterable<PluginProperty> properties, @Nullable final String pluginName,
                              final CallContext callContext, final InternalCallContext internalCallContext) throws PaymentApiException {
-        return run(initialState, isApiPayment, transactionType, account, paymentMethodId, directPaymentId, directPaymentExternalKey, directPaymentTransactionExternalKey,
+        return run(retrySMHelper.getInitialState(), isApiPayment, transactionType, account, paymentMethodId, directPaymentId, directPaymentExternalKey, directPaymentTransactionExternalKey,
                    amount, currency, properties, pluginName, callContext, internalCallContext);
     }
 
@@ -110,10 +101,10 @@ public class PluginControlledDirectPaymentAutomatonRunner extends DirectPaymentA
         try {
 
             final OperationCallback callback = createOperationCallback(transactionType, directPaymentStateContext);
-            final LeavingStateCallback leavingStateCallback = new RetryLeavingStateCallback(this, directPaymentStateContext, paymentDao, initialState, retriedState, transactionType);
+            final LeavingStateCallback leavingStateCallback = new RetryLeavingStateCallback(this, directPaymentStateContext, paymentDao, retrySMHelper.getInitialState(), retrySMHelper.getRetriedState(), transactionType);
             final EnteringStateCallback enteringStateCallback = new RetryEnteringStateCallback(this, directPaymentStateContext, retryServiceScheduler);
 
-            state.runOperation(retryOperation, callback, enteringStateCallback, leavingStateCallback);
+            state.runOperation(retrySMHelper.getRetryOperation(), callback, enteringStateCallback, leavingStateCallback);
 
         } catch (MissingEntryException e) {
             throw new PaymentApiException(e.getCause(), ErrorCode.PAYMENT_INTERNAL_ERROR, Objects.firstNonNull(e.getMessage(), ""));
@@ -137,7 +128,7 @@ public class PluginControlledDirectPaymentAutomatonRunner extends DirectPaymentA
             final LeavingStateCallback leavingStateCallback = new RetryNoopLeavingStateCallback();
             final EnteringStateCallback enteringStateCallback = new RetryEnteringStateCallback(this, paymentStateContext, retryServiceScheduler);
 
-            initialState.runOperation(retryOperation, callback, enteringStateCallback, leavingStateCallback);
+            retrySMHelper.getInitialState().runOperation(retrySMHelper.getRetryOperation(), callback, enteringStateCallback, leavingStateCallback);
 
         } catch (MissingEntryException e) {
             throw new PaymentApiException(e.getCause(), ErrorCode.PAYMENT_INTERNAL_ERROR, Objects.firstNonNull(e.getMessage(), ""));
@@ -153,24 +144,6 @@ public class PluginControlledDirectPaymentAutomatonRunner extends DirectPaymentA
         return paymentStateContext.getResult();
     }
 
-
-    // STEPH to be moved
-    public final State fetchState(final String stateName) {
-        try {
-            return retryStateMachine.getState(stateName);
-        } catch (MissingEntryException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
-    private final Operation fetchRetryOperation() {
-        try {
-            return retryStateMachine.getOperation("OP_RETRY");
-        } catch (MissingEntryException e) {
-            throw new RuntimeException(e);
-        }
-    }
-
     @VisibleForTesting
     RetryableDirectPaymentStateContext createContext(final boolean isApiPayment, final TransactionType transactionType, final Account account, @Nullable final UUID paymentMethodId,
                                                      @Nullable final UUID directPaymentId, @Nullable final String directPaymentExternalKey, final String directPaymentTransactionExternalKey,
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryStateMachineHelper.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryStateMachineHelper.java
new file mode 100644
index 0000000..30bb1ed
--- /dev/null
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryStateMachineHelper.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2014 Groupon, Inc
+ * Copyright 2014 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.payment.core.sm;
+
+import javax.inject.Inject;
+import javax.inject.Named;
+
+import org.killbill.automaton.MissingEntryException;
+import org.killbill.automaton.Operation;
+import org.killbill.automaton.State;
+import org.killbill.automaton.StateMachine;
+import org.killbill.automaton.StateMachineConfig;
+import org.killbill.billing.payment.glue.PaymentModule;
+
+public class RetryStateMachineHelper {
+
+    /**
+     * Those need to match RetryStates.xml
+     */
+    private static final String RETRY_STATE_MACHINE_NAME = "PAYMENT_RETRY";
+    private final String RETRY_OPERATION_NAME = "OP_RETRY";
+    private static final String INIT_STATE_NAME = "INIT";
+    private static final String RETRIED_STATE_NAME = "RETRIED";
+
+    private final StateMachineConfig retryStateMachineConfig;
+    private final StateMachine retryStateMachine;
+    private final Operation retryOperation;
+    private final State initialState;
+    private final State retriedState;
+
+    @Inject
+    public RetryStateMachineHelper(@Named(PaymentModule.STATE_MACHINE_RETRY) final StateMachineConfig retryStateMachineConfig) throws MissingEntryException {
+        this.retryStateMachineConfig = retryStateMachineConfig;
+        this.retryStateMachine = retryStateMachineConfig.getStateMachine(RETRY_STATE_MACHINE_NAME);
+        this.retryOperation = retryStateMachine.getOperation(RETRY_OPERATION_NAME);
+        this.initialState = retryStateMachine.getState(INIT_STATE_NAME);
+        this.retriedState = retryStateMachine.getState(RETRIED_STATE_NAME);
+    }
+
+    public State getState(final String stateName) throws MissingEntryException {
+        return retryStateMachine.getState(stateName);
+    }
+
+    public StateMachineConfig getRetryStateMachineConfig() {
+        return retryStateMachineConfig;
+    }
+
+    public StateMachine getRetryStateMachine() {
+        return retryStateMachine;
+    }
+
+    public Operation getRetryOperation() {
+        return retryOperation;
+    }
+
+    public State getInitialState() {
+        return initialState;
+    }
+
+    public State getRetriedState() {
+        return retriedState;
+    }
+}
diff --git a/payment/src/main/java/org/killbill/billing/payment/glue/PaymentModule.java b/payment/src/main/java/org/killbill/billing/payment/glue/PaymentModule.java
index 4074131..18ca713 100644
--- a/payment/src/main/java/org/killbill/billing/payment/glue/PaymentModule.java
+++ b/payment/src/main/java/org/killbill/billing/payment/glue/PaymentModule.java
@@ -41,7 +41,9 @@ import org.killbill.billing.payment.core.Janitor;
 import org.killbill.billing.payment.core.PaymentGatewayProcessor;
 import org.killbill.billing.payment.core.PaymentMethodProcessor;
 import org.killbill.billing.payment.core.PluginControlledPaymentProcessor;
+import org.killbill.billing.payment.core.sm.PaymentStateMachineHelper;
 import org.killbill.billing.payment.core.sm.PluginControlledDirectPaymentAutomatonRunner;
+import org.killbill.billing.payment.core.sm.RetryStateMachineHelper;
 import org.killbill.billing.payment.dao.DefaultPaymentDao;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.plugin.api.PaymentPluginApi;
@@ -91,6 +93,7 @@ public class PaymentModule extends KillBillModule {
 
         bind(Janitor.class).asEagerSingleton();
     }
+
     protected void installRetryEngines() {
         bind(DefaultRetryService.class).asEagerSingleton();
         bind(RetryService.class).annotatedWith(Names.named(RETRYABLE_NAMED)).to(DefaultRetryService.class);
@@ -103,9 +106,11 @@ public class PaymentModule extends KillBillModule {
 
         bind(StateMachineProvider.class).annotatedWith(Names.named(STATE_MACHINE_RETRY)).toInstance(new StateMachineProvider("org/killbill/billing/payment/retry/RetryStates.xml"));
         bind(StateMachineConfig.class).annotatedWith(Names.named(STATE_MACHINE_RETRY)).toProvider(Key.get(StateMachineProvider.class, Names.named(STATE_MACHINE_RETRY)));
+        bind(RetryStateMachineHelper.class).asEagerSingleton();
 
         bind(StateMachineProvider.class).annotatedWith(Names.named(STATE_MACHINE_PAYMENT)).toInstance(new StateMachineProvider("org/killbill/billing/payment/PaymentStates.xml"));
         bind(StateMachineConfig.class).annotatedWith(Names.named(STATE_MACHINE_PAYMENT)).toProvider(Key.get(StateMachineProvider.class, Names.named(STATE_MACHINE_PAYMENT)));
+        bind(PaymentStateMachineHelper.class).asEagerSingleton();
     }
 
     public static final class StateMachineProvider implements Provider<StateMachineConfig> {
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryableDirectPaymentAutomatonRunner.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryableDirectPaymentAutomatonRunner.java
index efcbaaa..ff5ae5b 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryableDirectPaymentAutomatonRunner.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryableDirectPaymentAutomatonRunner.java
@@ -56,8 +56,10 @@ public class MockRetryableDirectPaymentAutomatonRunner extends PluginControlledD
     private RetryableDirectPaymentStateContext context;
 
     @Inject
-    public MockRetryableDirectPaymentAutomatonRunner(@Named(PaymentModule.STATE_MACHINE_PAYMENT) final StateMachineConfig stateMachineConfig, @Named(PaymentModule.STATE_MACHINE_RETRY) final StateMachineConfig retryStateMachine, final PaymentDao paymentDao, final GlobalLocker locker, final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry, final OSGIServiceRegistration<PaymentControlPluginApi> retryPluginRegistry, final Clock clock, final TagInternalApi tagApi, final DirectPaymentProcessor directPaymentProcessor, @Named(RETRYABLE_NAMED) final RetryServiceScheduler retryServiceScheduler, final PaymentConfig paymentConfig, @com.google.inject.name.Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor) {
-        super(stateMachineConfig, retryStateMachine, paymentDao, locker, pluginRegistry, retryPluginRegistry, clock, tagApi, directPaymentProcessor, retryServiceScheduler, paymentConfig, executor);
+    public MockRetryableDirectPaymentAutomatonRunner(@Named(PaymentModule.STATE_MACHINE_PAYMENT) final StateMachineConfig stateMachineConfig, @Named(PaymentModule.STATE_MACHINE_RETRY) final StateMachineConfig retryStateMachine, final PaymentDao paymentDao, final GlobalLocker locker, final OSGIServiceRegistration<PaymentPluginApi> pluginRegistry, final OSGIServiceRegistration<PaymentControlPluginApi> retryPluginRegistry, final Clock clock, final TagInternalApi tagApi, final DirectPaymentProcessor directPaymentProcessor,
+                                                     @Named(RETRYABLE_NAMED) final RetryServiceScheduler retryServiceScheduler, final PaymentConfig paymentConfig, @com.google.inject.name.Named(PLUGIN_EXECUTOR_NAMED) final ExecutorService executor,
+                                                     final PaymentStateMachineHelper paymentSMHelper, final RetryStateMachineHelper retrySMHelper) {
+        super(stateMachineConfig, paymentDao, locker, pluginRegistry, retryPluginRegistry, clock, directPaymentProcessor, retryServiceScheduler, paymentConfig, executor, paymentSMHelper, retrySMHelper);
     }
 
     @Override
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentAutomatonDAOHelper.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentAutomatonDAOHelper.java
index a026985..0049c75 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentAutomatonDAOHelper.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentAutomatonDAOHelper.java
@@ -139,6 +139,6 @@ public class TestDirectPaymentAutomatonDAOHelper extends PaymentTestSuiteWithEmb
                                                                   internalCallContext,
                                                                   callContext);
 
-        return new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext);
+        return new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext, paymentSMHelper);
     }
 }
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentEnteringStateCallback.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentEnteringStateCallback.java
index b976bc7..a96a69e 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentEnteringStateCallback.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentEnteringStateCallback.java
@@ -68,7 +68,7 @@ public class TestDirectPaymentEnteringStateCallback extends PaymentTestSuiteWith
                                                                   ImmutableList.<PluginProperty>of(),
                                                                   internalCallContext,
                                                                   callContext);
-        daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext);
+        daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext, paymentSMHelper);
         callback = new DirectPaymentEnteringStateTestCallback(daoHelper, directPaymentStateContext);
 
         Mockito.when(state.getName()).thenReturn("NEW_STATE");
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentLeavingStateCallback.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentLeavingStateCallback.java
index 3e0e11c..5441999 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentLeavingStateCallback.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentLeavingStateCallback.java
@@ -134,7 +134,7 @@ public class TestDirectPaymentLeavingStateCallback extends PaymentTestSuiteWithE
             paymentDao.insertDirectPaymentWithFirstTransaction(newPaymentModelDao, newPaymentTransactionModelDao, internalCallContext);
         }
 
-        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext);
+        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext, paymentSMHelper);
         callback = new DirectPaymentLeavingStateTestCallback(daoHelper);
 
         Mockito.when(state.getName()).thenReturn("NEW_STATE");
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
index 0dbb4ff..111f6f7 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
@@ -122,7 +122,7 @@ public class TestDirectPaymentOperation extends PaymentTestSuiteNoDB {
         final PaymentDao paymentDao = Mockito.mock(PaymentDao.class);
         Mockito.when(paymentDao.getPaymentMethodIncludedDeleted(directPaymentStateContext.getPaymentMethodId(), internalCallContext)).thenReturn(paymentMethodModelDao);
 
-        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext);
+        final DirectPaymentAutomatonDAOHelper daoHelper = new DirectPaymentAutomatonDAOHelper(directPaymentStateContext, clock.getUTCNow(), paymentDao, registry, internalCallContext, paymentSMHelper);
         directPaymentOperation = new DirectPaymentOperationTest(paymentPluginStatus, daoHelper, locker, paymentPluginDispatcher, directPaymentStateContext);
     }
 
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestRetryableDirectPayment.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestRetryableDirectPayment.java
index 42bc3b4..0fc3f10 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestRetryableDirectPayment.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestRetryableDirectPayment.java
@@ -101,6 +101,10 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
     @Inject
     @Named(PLUGIN_EXECUTOR_NAMED)
     private ExecutorService executor;
+    @Inject
+    private PaymentStateMachineHelper paymentSMHelper;
+    @Inject
+    private RetryStateMachineHelper retrySMHelper;
 
     private Account account;
     private DateTime utcNow;
@@ -157,7 +161,9 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
                 directPaymentProcessor,
                 retryServiceScheduler,
                 paymentConfig,
-                executor);
+                executor,
+                paymentSMHelper,
+                retrySMHelper);
 
         directPaymentStateContext =
                 new RetryableDirectPaymentStateContext(MockPaymentControlProviderPlugin.PLUGIN_NAME,
@@ -193,6 +199,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
                                                          locker,
                                                          executor,
                                                          runner,
+                                                         retrySMHelper,
                                                          clock);
 
     }
@@ -461,7 +468,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
         runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
               .setContext(directPaymentStateContext);
 
-        final State state = runner.fetchState("RETRIED");
+        final State state = retrySMHelper.getRetriedState();
         final UUID directTransactionId = UUID.randomUUID();
         paymentDao.insertPaymentAttemptWithProperties(new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
                                                                                  directPaymentExternalKey, directTransactionId, directPaymentTransactionExternalKey,
@@ -509,7 +516,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
         runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
               .setContext(directPaymentStateContext);
 
-        final State state = runner.fetchState("RETRIED");
+        final State state = retrySMHelper.getRetriedState();
         final UUID directTransactionId = UUID.randomUUID();
         paymentDao.insertPaymentAttemptWithProperties(new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
                                                                                  directPaymentExternalKey, directTransactionId, directPaymentTransactionExternalKey,
@@ -556,7 +563,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
         runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
               .setContext(directPaymentStateContext);
 
-        final State state = runner.fetchState("RETRIED");
+        final State state = retrySMHelper.getRetriedState();
         final UUID directTransactionId = UUID.randomUUID();
         paymentDao.insertPaymentAttemptWithProperties(new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
                                                                                  directPaymentExternalKey, directTransactionId, directPaymentTransactionExternalKey,
@@ -611,7 +618,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
         runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
               .setContext(directPaymentStateContext);
 
-        final State state = runner.fetchState("RETRIED");
+        final State state = retrySMHelper.getRetriedState();
         final UUID directTransactionId = UUID.randomUUID();
         final UUID directPaymentId = UUID.randomUUID();
         final PaymentAttemptModelDao attempt = new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
@@ -653,7 +660,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
         runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
               .setContext(directPaymentStateContext);
 
-        final State state = runner.fetchState("RETRIED");
+        final State state = retrySMHelper.getRetriedState();
         final UUID directTransactionId = UUID.randomUUID();
         final UUID directPaymentId = UUID.randomUUID();
         final PaymentAttemptModelDao attempt = new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
@@ -701,7 +708,7 @@ public class TestRetryableDirectPayment extends PaymentTestSuiteNoDB {
             runner.setOperationCallback(mockRetryAuthorizeOperationCallback)
                   .setContext(directPaymentStateContext);
 
-            final State state = runner.fetchState("RETRIED");
+            final State state = retrySMHelper.getRetriedState();
             final UUID directTransactionId = UUID.randomUUID();
             final UUID directPaymentId = UUID.randomUUID();
             final PaymentAttemptModelDao attempt = new PaymentAttemptModelDao(account.getId(), paymentMethodId, utcNow, utcNow,
diff --git a/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteNoDB.java b/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteNoDB.java
index 6e34088..485d4d4 100644
--- a/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteNoDB.java
+++ b/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteNoDB.java
@@ -26,6 +26,7 @@ import org.killbill.billing.payment.api.DirectPaymentApi;
 import org.killbill.billing.payment.core.DirectPaymentProcessor;
 import org.killbill.billing.payment.core.PaymentMethodProcessor;
 import org.killbill.billing.payment.core.PluginControlledPaymentProcessor;
+import org.killbill.billing.payment.core.sm.PaymentStateMachineHelper;
 import org.killbill.billing.payment.core.sm.PluginControlledDirectPaymentAutomatonRunner;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.glue.TestPaymentModuleNoDB;
@@ -65,6 +66,8 @@ public abstract class PaymentTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
     @Inject
     protected PaymentDao paymentDao;
     @Inject
+    protected PaymentStateMachineHelper paymentSMHelper;
+    @Inject
     protected DirectPaymentProcessor paymentProcessor;
     @Inject
     protected PluginControlledPaymentProcessor pluginControlledPaymentProcessor;
diff --git a/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java b/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java
index be1b0ef..ad00182 100644
--- a/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java
+++ b/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java
@@ -25,6 +25,7 @@ import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.DirectPaymentApi;
 import org.killbill.billing.payment.core.DirectPaymentProcessor;
 import org.killbill.billing.payment.core.PaymentMethodProcessor;
+import org.killbill.billing.payment.core.sm.PaymentStateMachineHelper;
 import org.killbill.billing.payment.dao.PaymentDao;
 import org.killbill.billing.payment.glue.TestPaymentModuleWithEmbeddedDB;
 import org.killbill.billing.payment.plugin.api.PaymentPluginApi;
@@ -60,6 +61,8 @@ public abstract class PaymentTestSuiteWithEmbeddedDB extends GuicyKillbillTestSu
     @Inject
     protected AccountInternalApi accountApi;
     @Inject
+    protected PaymentStateMachineHelper paymentSMHelper;
+    @Inject
     protected PaymentDao paymentDao;
     @Inject
     protected TestPaymentHelper testHelper;