killbill-memoizeit

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 4587b28..2010258 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
@@ -222,8 +222,6 @@ public class PaymentAutomatonRunner {
                 throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR, Objects.firstNonNull(e.getMessage(), ""));
             } else if (e.getCause() instanceof PaymentApiException) {
                 throw (PaymentApiException) e.getCause();
-            } else if (e.getCause() instanceof TimeoutException) {
-                throw new PaymentApiException(e.getCause(), ErrorCode.PAYMENT_PLUGIN_TIMEOUT, accountId, Objects.firstNonNull(e.getMessage(), ""));
             } else {
                 throw new PaymentApiException(e.getCause(), ErrorCode.PAYMENT_INTERNAL_ERROR, Objects.firstNonNull(e.getMessage(), ""));
             }

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 c8db8e8..bfb1c21 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
@@ -112,6 +112,9 @@ public abstract class PaymentOperation extends OperationCallbackBase<PaymentTran
                                                                                                 null,
                                                                                                 null);
         paymentStateContext.setPaymentTransactionInfoPlugin(paymentInfoPlugin);
+        if (e.getCause() instanceof OperationException) {
+            return (OperationException) e.getCause();
+        }
         if (e instanceof OperationException) {
             return (OperationException) 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 c64ef5f..f1c4a17 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
@@ -20,10 +20,10 @@ package org.killbill.billing.payment.api;
 
 import java.math.BigDecimal;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.UUID;
-import java.util.concurrent.TimeoutException;
 
 import javax.annotation.Nullable;
 
@@ -35,13 +35,16 @@ import org.killbill.billing.control.plugin.api.PaymentControlApiException;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
 import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.osgi.api.OSGIServiceDescriptor;
 import org.killbill.billing.payment.MockRecurringInvoiceItem;
 import org.killbill.billing.payment.PaymentTestSuiteWithEmbeddedDB;
 import org.killbill.billing.payment.dao.PaymentAttemptModelDao;
 import org.killbill.billing.payment.dao.PaymentSqlDao;
 import org.killbill.billing.payment.invoice.InvoicePaymentControlPluginApi;
+import org.killbill.billing.payment.plugin.api.PaymentPluginApiException;
 import org.killbill.billing.payment.plugin.api.PaymentPluginStatus;
 import org.killbill.billing.payment.provider.ExternalPaymentProviderPlugin;
+import org.killbill.billing.payment.provider.MockPaymentControlProviderPlugin;
 import org.killbill.billing.payment.provider.MockPaymentProviderPlugin;
 import org.killbill.billing.util.entity.Pagination;
 import org.killbill.bus.api.PersistentBus.EventBusException;
@@ -63,6 +66,8 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
 
     private MockPaymentProviderPlugin mockPaymentProviderPlugin;
 
+    private MockPaymentControlProviderPlugin mockPaymentControlProviderPlugin;
+
     final PaymentOptions INVOICE_PAYMENT = new PaymentOptions() {
         @Override
         public boolean isExternalPayment() {
@@ -75,6 +80,18 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
         }
     };
 
+    final PaymentOptions CONTROL_PLUGIN_OPTIONS = new PaymentOptions() {
+        @Override
+        public boolean isExternalPayment() {
+            return true;
+        }
+
+        @Override
+        public List<String> getPaymentControlPluginNames() {
+            return Arrays.asList(MockPaymentControlProviderPlugin.PLUGIN_NAME);
+        }
+    };
+
     private Account account;
 
     @BeforeClass(groups = "slow")
@@ -88,6 +105,24 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
         super.beforeMethod();
         mockPaymentProviderPlugin.clear();
         account = testHelper.createTestAccount("bobo@gmail.com", true);
+
+        mockPaymentControlProviderPlugin = new MockPaymentControlProviderPlugin();;
+        controlPluginRegistry.registerService(new OSGIServiceDescriptor() {
+            @Override
+            public String getPluginSymbolicName() {
+                return null;
+            }
+
+            @Override
+            public String getPluginName() {
+                return MockPaymentControlProviderPlugin.PLUGIN_NAME;
+            }
+
+            @Override
+            public String getRegistrationName() {
+                return MockPaymentControlProviderPlugin.PLUGIN_NAME;
+            }
+        }, mockPaymentControlProviderPlugin);
     }
 
     @Test(groups = "slow")
@@ -250,6 +285,55 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
 
     }
 
+    @Test(groups = "slow")
+    public void testCreatePurchasePaymentPluginException() {
+        mockPaymentProviderPlugin.makeNextPaymentFailWithException();
+
+        final BigDecimal requestedAmount = BigDecimal.TEN;
+        final String paymentExternalKey = "pay external key";
+        final String transactionExternalKey = "txn external key";
+        try {
+            paymentApi.createPurchase(account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED,
+                                      paymentExternalKey, transactionExternalKey, ImmutableList.<PluginProperty>of(), callContext);
+            fail();
+        } catch (PaymentApiException e) {
+            assertTrue(e.getCause() instanceof PaymentPluginApiException);
+        }
+    }
+
+    @Test(groups = "slow")
+    public void testCreatePurchaseWithControlPaymentPluginException() throws Exception {
+        mockPaymentProviderPlugin.makeNextPaymentFailWithException();
+
+        final BigDecimal requestedAmount = BigDecimal.TEN;
+        final String paymentExternalKey = "pay controle external key";;
+        final String transactionExternalKey = "txn control external key";
+        try {
+            paymentApi.createPurchaseWithPaymentControl(
+                    account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED,
+                    paymentExternalKey, transactionExternalKey, ImmutableList.<PluginProperty>of(), CONTROL_PLUGIN_OPTIONS, callContext);
+            fail();
+        } catch (PaymentApiException e) {
+            assertTrue(e.getCause() instanceof PaymentPluginApiException);
+        }
+    }
+
+    @Test(groups = "slow")
+    public void testCreatePurchaseWithControlPluginException() throws Exception {
+        mockPaymentControlProviderPlugin.throwsException(new IllegalStateException());
+
+        final BigDecimal requestedAmount = BigDecimal.TEN;
+        final String paymentExternalKey = "pay controle external key";;
+        final String transactionExternalKey = "txn control external key";
+        try {
+            paymentApi.createPurchaseWithPaymentControl(
+                    account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED,
+                    paymentExternalKey, transactionExternalKey, ImmutableList.<PluginProperty>of(), CONTROL_PLUGIN_OPTIONS, callContext);
+            fail();
+        } catch (PaymentApiException e) {
+            assertTrue(e.getCause() instanceof IllegalStateException);
+        }
+    }
 
     @Test(groups = "slow")
     public void testCreateSuccessAuthVoid() throws PaymentApiException {
@@ -1180,8 +1264,25 @@ public class TestPaymentApi extends PaymentTestSuiteWithEmbeddedDB {
 
         mockPaymentProviderPlugin.makePluginWaitSomeMilliseconds((int) (paymentConfig.getPaymentPluginTimeout().getMillis() + 100));
         try {
-            final Payment payment = paymentApi.createPurchase(account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED, paymentExternalKey, transactionExternalKey,
-                                                              ImmutableList.<PluginProperty>of(), callContext);
+            paymentApi.createPurchase(account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED,
+                                      paymentExternalKey, transactionExternalKey, ImmutableList.<PluginProperty>of(), callContext);
+            fail();
+        } catch (PaymentApiException e) {
+            assertEquals(e.getCode(), ErrorCode.PAYMENT_PLUGIN_TIMEOUT.getCode());
+        }
+    }
+
+    @Test(groups = "slow")
+    public void testCreatePurchaseWithControlTimeout() throws Exception {
+        final BigDecimal requestedAmount = BigDecimal.ONE;
+        final String paymentExternalKey = "111111";
+        final String transactionExternalKey = "11111";
+
+        mockPaymentProviderPlugin.makePluginWaitSomeMilliseconds((int) (paymentConfig.getPaymentPluginTimeout().getMillis() + 100));
+        try {
+            paymentApi.createPurchaseWithPaymentControl(
+                    account, account.getPaymentMethodId(), null, requestedAmount, Currency.AED, paymentExternalKey,
+                    transactionExternalKey, ImmutableList.<PluginProperty>of(), CONTROL_PLUGIN_OPTIONS, callContext);
             fail();
         } catch (PaymentApiException e) {
             assertEquals(e.getCode(), ErrorCode.PAYMENT_PLUGIN_TIMEOUT.getCode());

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 dc7b36f..786dae5 100644
--- a/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java
+++ b/payment/src/test/java/org/killbill/billing/payment/PaymentTestSuiteWithEmbeddedDB.java
@@ -20,6 +20,7 @@ package org.killbill.billing.payment;
 
 import org.killbill.billing.GuicyKillbillTestSuiteWithEmbeddedDB;
 import org.killbill.billing.account.api.AccountInternalApi;
+import org.killbill.billing.control.plugin.api.PaymentControlPluginApi;
 import org.killbill.billing.invoice.api.InvoiceInternalApi;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.payment.api.PaymentApi;
@@ -59,6 +60,8 @@ public abstract class PaymentTestSuiteWithEmbeddedDB extends GuicyKillbillTestSu
     @Inject
     protected OSGIServiceRegistration<PaymentPluginApi> registry;
     @Inject
+    protected OSGIServiceRegistration<PaymentControlPluginApi> controlPluginRegistry;
+    @Inject
     protected PersistentBus eventBus;
     @Inject
     protected PaymentApi paymentApi;

diff --git a/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentControlProviderPlugin.java b/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentControlProviderPlugin.java
index 0327309..c150594 100644
--- a/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentControlProviderPlugin.java
+++ b/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentControlProviderPlugin.java
@@ -35,6 +35,7 @@ public class MockPaymentControlProviderPlugin implements PaymentControlPluginApi
 
     private boolean isAborted;
     private DateTime nextRetryDate;
+    private Exception exception;
 
     public MockPaymentControlProviderPlugin setAborted(final boolean isAborted) {
         this.isAborted = isAborted;
@@ -46,8 +47,23 @@ public class MockPaymentControlProviderPlugin implements PaymentControlPluginApi
         return this;
     }
 
+    public MockPaymentControlProviderPlugin throwsException(PaymentControlApiException exception) {
+        this.exception = exception;
+        return this;
+    }
+
+    public MockPaymentControlProviderPlugin throwsException(RuntimeException exception) {
+        this.exception = exception;
+        return this;
+    }
+
     @Override
     public PriorPaymentControlResult priorCall(final PaymentControlContext paymentControlContext, final Iterable<PluginProperty> properties) throws PaymentControlApiException {
+        if (exception instanceof PaymentControlApiException) {
+            throw (PaymentControlApiException) exception;
+        } else if (exception instanceof RuntimeException) {
+            throw (RuntimeException) exception;
+        }
         return new DefaultPriorPaymentControlResult(isAborted);
     }