killbill-aplcache

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/OperationControlCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/OperationControlCallback.java
index d0d61d2..22b0d54 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/OperationControlCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/OperationControlCallback.java
@@ -157,7 +157,7 @@ public abstract class OperationControlCallback extends OperationCallbackBase<Pay
     }
 
     @Override
-    protected OperationException unwrapExceptionFromDispatchedTask(final PaymentStateContext paymentStateContext, final Exception e) {
+    protected OperationException unwrapExceptionFromDispatchedTask(final Exception e) {
         // If this is an ExecutionException we attempt to extract the cause first
         final Throwable originalExceptionOrCausePossiblyOperationException = e instanceof ExecutionException ? MoreObjects.firstNonNull(e.getCause(), e) : e;
 

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/OperationCallbackBase.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/OperationCallbackBase.java
index 34341d9..1f52802 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/OperationCallbackBase.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/OperationCallbackBase.java
@@ -72,14 +72,14 @@ public abstract class OperationCallbackBase<CallbackOperationResult, CallbackOpe
             logger.debug("Successful plugin(s) call of {} for account {} with result {}", pluginNames, account.getExternalKey(), operationResult);
             return operationResult;
         } catch (final ExecutionException e) {
-            throw unwrapExceptionFromDispatchedTask(paymentStateContext, e);
+            throw unwrapExceptionFromDispatchedTask(e);
         } catch (final TimeoutException e) {
             logger.warn("TimeoutException while executing the plugin(s) {}", pluginNames);
-            throw unwrapExceptionFromDispatchedTask(paymentStateContext, e);
+            throw unwrapExceptionFromDispatchedTask(e);
         } catch (final InterruptedException e) {
             Thread.currentThread().interrupt();
             logger.warn("InterruptedException while executing the following plugin(s): {}", pluginNames);
-            throw unwrapExceptionFromDispatchedTask(paymentStateContext, e);
+            throw unwrapExceptionFromDispatchedTask(e);
         }
     }
 
@@ -91,5 +91,5 @@ public abstract class OperationCallbackBase<CallbackOperationResult, CallbackOpe
     //
     protected abstract CallbackOperationResult doCallSpecificOperationCallback() throws CallbackOperationException;
 
-    protected abstract OperationException unwrapExceptionFromDispatchedTask(final PaymentStateContext paymentStateContext, final Exception e);
+    protected abstract OperationException unwrapExceptionFromDispatchedTask(final Exception e);
 }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
index 9388e2f..e269932 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
@@ -72,22 +72,29 @@ public abstract class PaymentOperation extends OperationCallbackBase<PaymentTran
 
     @Override
     public OperationResult doOperationCallback() throws OperationException {
+        final String pluginName;
         try {
-            final String pluginName = daoHelper.getPaymentProviderPluginName();
+            pluginName = daoHelper.getPaymentProviderPluginName();
             this.plugin = daoHelper.getPaymentPluginApi(pluginName);
+        } catch (final PaymentApiException e) {
+            throw convertToUnknownTransactionStatusAndErroredPaymentState(e);
+        }
 
-            if (paymentStateContext.shouldLockAccountAndDispatch()) {
-                return doOperationCallbackWithDispatchAndAccountLock(pluginName);
-            } else {
+        if (paymentStateContext.shouldLockAccountAndDispatch()) {
+            // This will already call unwrapExceptionFromDispatchedTask
+            return doOperationCallbackWithDispatchAndAccountLock(pluginName);
+        } else {
+            try {
                 return doSimpleOperationCallback();
+            } catch (final Exception e) {
+                // We need to unwrap OperationException (see doSimpleOperationCallback below)
+                throw unwrapExceptionFromDispatchedTask(e);
             }
-        } catch (final Exception e) {
-            throw convertToUnknownTransactionStatusAndErroredPaymentState(e);
         }
     }
 
     @Override
-    protected OperationException unwrapExceptionFromDispatchedTask(final PaymentStateContext paymentStateContext, final Exception e) {
+    protected OperationException unwrapExceptionFromDispatchedTask(final Exception e) {
         // If this is an ExecutionException we attempt to extract the cause first
         final Throwable originalExceptionOrCausePossiblyOperationException = e instanceof ExecutionException ? MoreObjects.firstNonNull(e.getCause(), e) : e;
 

diff --git a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
index a3cca2f..d2c8de8 100644
--- a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
+++ b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApi.java
@@ -1270,11 +1270,7 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
 
                 assertNotNull(thrownException);
 
-                Throwable operationException = thrownException.getCause();
-                assertNotNull(operationException);
-                assertTrue(operationException instanceof OperationException);
-
-                Throwable timeoutException = operationException.getCause();
+                Throwable timeoutException = thrownException.getCause();
                 assertNotNull(timeoutException);
                 assertTrue(timeoutException instanceof TimeoutException);