diff --git a/modules/payment_form/tests/src/Unit/Entity/Payment/PaymentFormUnitTest.php b/modules/payment_form/tests/src/Unit/Entity/Payment/PaymentFormUnitTest.php
index f8e3208..70f142a 100644
--- a/modules/payment_form/tests/src/Unit/Entity/Payment/PaymentFormUnitTest.php
+++ b/modules/payment_form/tests/src/Unit/Entity/Payment/PaymentFormUnitTest.php
@@ -212,7 +212,12 @@ class PaymentFormUnitTest extends UnitTestCase {
    */
   public function testSubmitForm() {
     $redirect_url = new Url($this->randomMachineName());
-    $payment_method_execution_status = new Response($redirect_url);
+    $response = new Response($redirect_url);
+
+    $result = $this->getMock('\Drupal\payment\PaymentExecutionResultInterface');
+    $result->expects($this->atLeastOnce())
+      ->method('getCompletionResponse')
+      ->willReturn($response);
 
     $form = array(
       'payment_method' => array(
@@ -230,21 +235,21 @@ class PaymentFormUnitTest extends UnitTestCase {
 
     $payment_method = $this->getMock('\Drupal\payment\Plugin\Payment\Method\PaymentMethodInterface');
 
-    $this->paymentMethodSelector->expects($this->once())
+    $this->paymentMethodSelector->expects($this->atLeastOnce())
       ->method('getPaymentMethod')
       ->will($this->returnValue($payment_method));
-    $this->paymentMethodSelector->expects($this->once())
+    $this->paymentMethodSelector->expects($this->atLeastOnce())
       ->method('submitConfigurationForm')
       ->with($form['payment_method'], $form_state);
 
-    $this->payment->expects($this->once())
+    $this->payment->expects($this->atLeastOnce())
       ->method('setPaymentMethod')
       ->with($payment_method);
-    $this->payment->expects($this->once())
+    $this->payment->expects($this->atLeastOnce())
       ->method('save');
-    $this->payment->expects($this->once())
+    $this->payment->expects($this->atLeastOnce())
       ->method('execute')
-      ->willReturn($payment_method_execution_status);
+      ->willReturn($result);
 
     $this->form->submitForm($form, $form_state);
   }
diff --git a/modules/payment_reference/src/Element/PaymentReferenceBase.php b/modules/payment_reference/src/Element/PaymentReferenceBase.php
index 1c0d617..a3a9fc1 100644
--- a/modules/payment_reference/src/Element/PaymentReferenceBase.php
+++ b/modules/payment_reference/src/Element/PaymentReferenceBase.php
@@ -307,7 +307,9 @@ abstract class PaymentReferenceBase extends FormElement implements FormElementIn
   protected function buildRefreshButton(array $element, FormStateInterface $form_state) {
     $payment_method_selector = $this->getPaymentMethodSelector($element, $form_state);
     $class = array('payment_reference-refresh-button');
-    if (!$element['#default_value'] && $element['#available_payment_id'] && (!$payment_method_selector->getPaymentMethod() || !$payment_method_selector->getPaymentMethod()->getPaymentExecutionResult()->hasCompleted())) {
+    if (!$element['#default_value']
+      && $element['#available_payment_id']
+      && (!$payment_method_selector->getPaymentMethod() || !$payment_method_selector->getPaymentMethod()->getPaymentExecutionResult()->hasCompleted())) {
       $class[] = 'payment-reference-hidden';
     }
     $plugin_id = $this->getPluginId();
diff --git a/modules/payment_reference/src/Plugin/Payment/Type/PaymentReference.php b/modules/payment_reference/src/Plugin/Payment/Type/PaymentReference.php
index 8d0210b..6d23d26 100644
--- a/modules/payment_reference/src/Plugin/Payment/Type/PaymentReference.php
+++ b/modules/payment_reference/src/Plugin/Payment/Type/PaymentReference.php
@@ -13,6 +13,7 @@ use Drupal\Core\Session\AccountInterface;
 use Drupal\Core\StringTranslation\TranslationInterface;
 use Drupal\Core\Url;
 use Drupal\payment\Plugin\Payment\Type\PaymentTypeBase;
+use Drupal\payment\Response;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 use Symfony\Component\EventDispatcher\EventDispatcherInterface;
 use Symfony\Component\HttpFoundation\RedirectResponse;
@@ -98,11 +99,11 @@ class PaymentReference extends PaymentTypeBase implements ContainerFactoryPlugin
    * {@inheritdoc}
    */
   protected function doGetResumeContextResponse() {
-    return new RedirectResponse((new Url('payment_reference.resume_context', array(
+    return new Response(new Url('payment_reference.resume_context', array(
       'payment' => $this->getPayment()->id(),
     ), array(
       'absolute' => TRUE,
-    )))->toString());
+    )));
   }
 
   /**
diff --git a/modules/payment_reference/tests/src/Unit/Element/PaymentReferenceBaseUnitTest.php b/modules/payment_reference/tests/src/Unit/Element/PaymentReferenceBaseUnitTest.php
index 9491e8a..02d7330 100644
--- a/modules/payment_reference/tests/src/Unit/Element/PaymentReferenceBaseUnitTest.php
+++ b/modules/payment_reference/tests/src/Unit/Element/PaymentReferenceBaseUnitTest.php
@@ -84,13 +84,6 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
   protected $stringTranslation;
 
   /**
-   * The temporary payment storage.
-   *
-   * @var \Drupal\Core\KeyValueStore\KeyValueStoreExpirableInterface|\PHPUnit_Framework_MockObject_MockObject
-   */
-  protected $temporaryPaymentStorage;
-
-  /**
    * The url generator.
    *
    * @var \Drupal\Core\Routing\UrlGeneratorInterface|\PHPUnit_Framework_MockObject_MockObject
@@ -121,8 +114,6 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
 
     $this->stringTranslation = $this->getStringTranslationStub();
 
-    $this->temporaryPaymentStorage = $this->getMock('\Drupal\Core\KeyValueStore\KeyValueStoreExpirableInterface');
-
     $this->urlGenerator = $this->getMock('\Drupal\Core\Routing\UrlGeneratorInterface');
 
     $configuration = array();
@@ -136,9 +127,6 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
     $this->element->expects($this->any())
       ->method('getPaymentQueue')
       ->willReturn($this->paymentQueue);
-    $this->element->expects($this->any())
-      ->method('getTemporaryPaymentStorage')
-      ->willReturn($this->temporaryPaymentStorage);
   }
 
   /**
@@ -226,7 +214,7 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
 
     $this->element = $this->getMockBuilder('\Drupal\payment_reference\Element\PaymentReferenceBase')
       ->setConstructorArgs(array($configuration, $plugin_id, $this->pluginDefinition, $this->requestStack, $this->paymentStorage, $this->stringTranslation, $this->dateFormatter, $this->linkGenerator, $this->paymentMethodSelectorManager, new Random()))
-      ->setMethods(array('getEntityFormDisplay', 'getPaymentMethodSelector', 'getTemporaryPaymentStorageKey', 'setTemporaryPayment'))
+      ->setMethods(array('getEntityFormDisplay', 'getPaymentMethodSelector'))
       ->getMockForAbstractClass();
 
     $form = array(
@@ -347,7 +335,7 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
       ->willReturn($has_completed);
 
     $payment_method = $this->getMock('\Drupal\payment\Plugin\Payment\Method\PaymentMethodInterface');
-    $payment_method->expects($this->atLeastOnce())
+    $payment_method->expects($this->any())
       ->method('getPayment')
       ->willReturn($payment);
     $payment_method->expects($this->atLeastOnce())
@@ -550,15 +538,12 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
 
     $this->element = $this->getMockBuilder('\Drupal\payment_reference\Element\PaymentReferenceBase')
       ->setConstructorArgs(array($configuration, $plugin_id, $this->pluginDefinition, $this->requestStack, $this->paymentStorage, $this->stringTranslation, $this->dateFormatter, $this->linkGenerator, $this->paymentMethodSelectorManager, new Random()))
-      ->setMethods(array('getPaymentMethodSelector', 'getTemporaryPaymentStorageKey'))
+      ->setMethods(array('getPaymentMethodSelector'))
       ->getMockForAbstractClass();
     $this->element->expects($this->atLeastOnce())
       ->method('getPaymentMethodSelector')
       ->with($element, $form_state)
       ->willReturn($payment_method_selector);
-    $this->element->expects($this->never())
-      ->method('getTemporaryPaymentStorageKey')
-      ->with($element, $form_state);
 
 
     $method = new \ReflectionMethod($this->element, 'buildCompletePaymentLink');
@@ -597,15 +582,12 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
 
     $this->element = $this->getMockBuilder('\Drupal\payment_reference\Element\PaymentReferenceBase')
       ->setConstructorArgs(array($configuration, $plugin_id, $this->pluginDefinition, $this->requestStack, $this->paymentStorage, $this->stringTranslation, $this->dateFormatter, $this->linkGenerator, $this->paymentMethodSelectorManager, new Random()))
-      ->setMethods(array('getPaymentMethodSelector', 'getTemporaryPaymentStorageKey'))
+      ->setMethods(array('getPaymentMethodSelector'))
       ->getMockForAbstractClass();
     $this->element->expects($this->atLeastOnce())
       ->method('getPaymentMethodSelector')
       ->with($element, $form_state)
       ->willReturn($payment_method_selector);
-    $this->element->expects($this->atLeastOnce())
-      ->method('getTemporaryPaymentStorageKey')
-      ->with($element, $form_state);
 
 
     $method = new \ReflectionMethod($this->element, 'buildCompletePaymentLink');
@@ -721,9 +703,6 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
     $payment_method_selector_id = $this->randomMachineName();
 
     $payment_method_selector = $this->getMock('\Drupal\payment\Plugin\Payment\MethodSelector\PaymentMethodSelectorInterface');
-    $payment_method_selector->expects($this->atLeastOnce())
-      ->method('setAllowedPaymentMethodIds')
-      ->with($limit_allowed_payment_method_ids);
 
     $this->paymentMethodSelectorManager->expects($this->atLeastOnce())
       ->method('createInstance')
@@ -794,7 +773,7 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
 
     $this->element = $this->getMockBuilder('\Drupal\payment_reference\Element\PaymentReferenceBase')
       ->setConstructorArgs(array($configuration, $plugin_id, $this->pluginDefinition, $this->requestStack, $this->paymentStorage, $this->stringTranslation, $this->dateFormatter, $this->linkGenerator, $this->paymentMethodSelectorManager, new Random()))
-      ->setMethods(array('getEntityFormDisplay', 'getPaymentMethodSelector', 'hasTemporaryPayment'))
+      ->setMethods(array('getEntityFormDisplay', 'getPaymentMethodSelector'))
       ->getMockForAbstractClass();
     $this->element->expects($this->atLeastOnce())
       ->method('getEntityFormDisplay')
@@ -803,10 +782,6 @@ class PaymentReferenceBaseUnitTest extends UnitTestCase {
       ->method('getPaymentMethodSelector')
       ->with($element, $form_state)
       ->willReturn($payment_method_selector);
-    $this->element->expects($this->atLeastOnce())
-      ->method('hasTemporaryPayment')
-      ->with($element, $form_state)
-      ->willReturn(TRUE);
 
     $build = $method->invoke($this->element, $element, $form_state);
     $this->assertInternalType('array', $build);
diff --git a/modules/payment_reference/tests/src/Unit/Plugin/Payment/Type/PaymentReferenceUnitTest.php b/modules/payment_reference/tests/src/Unit/Plugin/Payment/Type/PaymentReferenceUnitTest.php
index b703778..8de72d5 100644
--- a/modules/payment_reference/tests/src/Unit/Plugin/Payment/Type/PaymentReferenceUnitTest.php
+++ b/modules/payment_reference/tests/src/Unit/Plugin/Payment/Type/PaymentReferenceUnitTest.php
@@ -233,16 +233,6 @@ class PaymentReferenceUnitTest extends UnitTestCase {
    * @covers ::doGetResumeContextResponse
    */
   public function testDoGetResumeContextResponse() {
-    $url_generator = $this->getMock('\Drupal\Core\Routing\UrlGeneratorInterface');
-    $url_generator->expects($this->atLeastOnce())
-      ->method('generateFromRoute')
-      ->willReturn($this->randomMachineName());
-
-    $container = new ContainerBuilder();
-    $container->set('url_generator', $url_generator);
-
-    \Drupal::setContainer($container);
-
     $response = $this->paymentType->getResumeContextResponse();
 
     $this->assertInstanceOf('\Drupal\payment\ResponseInterface', $response);