diff --git a/app/code/Magento/Braintree/Model/PaymentMethod.php b/app/code/Magento/Braintree/Model/PaymentMethod.php
index 3f85a3d79b563..4c263888bd4a8 100644
--- a/app/code/Magento/Braintree/Model/PaymentMethod.php
+++ b/app/code/Magento/Braintree/Model/PaymentMethod.php
@@ -269,6 +269,7 @@ public function assignData(\Magento\Framework\DataObject $data)
 
         $infoInstance->setAdditionalInformation('cc_last4', $additionalData->getData('cc_last4'));
         $infoInstance->setAdditionalInformation('cc_token', $additionalData->getData('cc_token'));
+        $infoInstance->setAdditionalInformation('store_in_vault', $additionalData->getData('store_in_vault'));
         $infoInstance->setAdditionalInformation(
             'payment_method_nonce',
             $additionalData->getData('payment_method_nonce')
diff --git a/app/code/Magento/Braintree/Test/Unit/Model/PaymentMethodTest.php b/app/code/Magento/Braintree/Test/Unit/Model/PaymentMethodTest.php
index 29e45a4016bca..1a6eaf854c9a0 100644
--- a/app/code/Magento/Braintree/Test/Unit/Model/PaymentMethodTest.php
+++ b/app/code/Magento/Braintree/Test/Unit/Model/PaymentMethodTest.php
@@ -276,17 +276,25 @@ public function testAssignData()
             ->with($ccExpYear)
             ->willReturnSelf();
 
-        $this->infoInstanceMock->expects($this->atLeastOnce())
+        $this->infoInstanceMock->expects($this->at(0))
             ->method('setAdditionalInformation')
-            ->willReturnMap(
-                [
-                    ['device_data', $deviceData],
-                    ['cc_last4', $ccLast4],
-                    ['cc_token', $ccToken],
-                    ['payment_method_nonce', $paymentMethodNonce],
-                    ['store_in_vault', $storeInVault]
-                ]
-            );
+            ->with('device_data', $deviceData);
+
+        $this->infoInstanceMock->expects($this->at(1))
+            ->method('setAdditionalInformation')
+            ->with('cc_last4', $ccLast4);
+
+        $this->infoInstanceMock->expects($this->at(2))
+            ->method('setAdditionalInformation')
+            ->with('cc_token', $ccToken);
+
+        $this->infoInstanceMock->expects($this->at(3))
+            ->method('setAdditionalInformation')
+            ->with('store_in_vault', $storeInVault);
+
+        $this->infoInstanceMock->expects($this->at(4))
+            ->method('setAdditionalInformation')
+            ->with('payment_method_nonce', $paymentMethodNonce);
 
         $this->model->assignData($data);
     }