killbill-memoizeit

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlCompleted.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlCompleted.java
index 6655210..0d3f97f 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlCompleted.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlCompleted.java
@@ -61,13 +61,12 @@ public class DefaultControlCompleted implements EnteringStateCallback {
                                    paymentStateContext.getCurrentTransaction().getId() :
                                    null;
 
+        // At this stage we can update the paymentAttempt state AND serialize the plugin properties. Control plugins will have had the opportunity to erase sensitive data if required.
+        retryablePaymentAutomatonRunner.getPaymentDao().updatePaymentAttemptWithProperties(attempt.getId(), transactionId, state.getName(), getSerializedProperties(), paymentStateContext.getInternalCallContext());
 
         if (retriedState.getName().equals(state.getName()) && !isUnknownTransaction()) {
-            retryablePaymentAutomatonRunner.getPaymentDao().updatePaymentAttemptWithProperties(attempt.getId(), transactionId, state.getName(), getSerializedProperties(), paymentStateContext.getInternalCallContext());
             retryServiceScheduler.scheduleRetry(ObjectType.PAYMENT_ATTEMPT, attempt.getId(), attempt.getId(), attempt.getTenantRecordId(),
                                                 paymentStateContext.getPaymentControlPluginNames(), paymentStateContext.getRetryDate());
-        } else {
-            retryablePaymentAutomatonRunner.getPaymentDao().updatePaymentAttempt(attempt.getId(), transactionId, state.getName(), paymentStateContext.getInternalCallContext());
         }
     }
 

diff --git a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApiWithControl.java b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApiWithControl.java
index b7d37a0..5d528a1 100644
--- a/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApiWithControl.java
+++ b/payment/src/test/java/org/killbill/billing/payment/api/TestPaymentApiWithControl.java
@@ -34,6 +34,8 @@ import org.killbill.billing.control.plugin.api.PaymentControlApiException;
 import org.killbill.billing.control.plugin.api.PaymentControlContext;
 import org.killbill.billing.control.plugin.api.PaymentControlPluginApi;
 import org.killbill.billing.control.plugin.api.PriorPaymentControlResult;
+import org.killbill.billing.payment.retry.DefaultFailureCallResult;
+import org.killbill.billing.payment.retry.DefaultOnSuccessPaymentControlResult;
 import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -112,12 +114,12 @@ public class TestPaymentApiWithControl extends PaymentTestSuiteWithEmbeddedDB {
 
         @Override
         public OnSuccessPaymentControlResult onSuccessCall(final PaymentControlContext context, final Iterable<PluginProperty> properties) throws PaymentControlApiException {
-            return null;
+            return new DefaultOnSuccessPaymentControlResult();
         }
 
         @Override
         public OnFailurePaymentControlResult onFailureCall(final PaymentControlContext context, final Iterable<PluginProperty> properties) throws PaymentControlApiException {
-            return null;
+            return new DefaultFailureCallResult(null);
         }
     }