MC-13644: Default billing and default shipping address lose their status after edit

Author: Joan He

app/code/Magento/Customer/Controller/Address/FormPost.php

@@ -120,8 +120,18 @@ protected function _extractAddress()
             \Magento\Customer\Api\Data\AddressInterface::class
         );
         $addressDataObject->setCustomerId($this->_getSession()->getCustomerId())
-            ->setIsDefaultBilling($this->getRequest()->getParam('default_billing', false))
-            ->setIsDefaultShipping($this->getRequest()->getParam('default_shipping', false));
+            ->setIsDefaultBilling(
+                $this->getRequest()->getParam(
+                    'default_billing',
+                    isset($existingAddressData['default_billing']) ? $existingAddressData['default_billing'] : false
+                )
+            )
+            ->setIsDefaultShipping(
+                $this->getRequest()->getParam(
+                    'default_shipping',
+                    isset($existingAddressData['default_shipping']) ? $existingAddressData['default_shipping'] : false
+                )
+            );
 
         return $addressDataObject;
     }

app/code/Magento/Customer/Test/Mftf/Test/StorefrontUpdateCustomerAddressTest.xml

@@ -18,9 +18,6 @@
             <testCaseId value="MAGETWO-97501"/>
             <group value="customer"/>
             <group value="update"/>
-            <skip>
-                <issueId value="MAGETWO-97504"/>
-            </skip>
         </annotations>
         <before>
             <createData entity="Simple_US_Customer_With_Different_Billing_Shipping_Addresses" stepKey="createCustomer"/>

app/code/Magento/Customer/Test/Unit/Controller/Address/FormPostTest.php

@@ -455,14 +455,20 @@ public function testExecute(
         $regionCode,
         $newRegionId,
         $newRegion,
-        $newRegionCode
+        $newRegionCode,
+        $existingDefaultBilling = false,
+        $existingDefaultShipping = false,
+        $setDefaultBilling = false,
+        $setDefaultShipping = false
     ): void {
         $existingAddressData = [
             'country_id' => $countryId,
             'region_id' => $regionId,
             'region' => $region,
             'region_code' => $regionCode,
-            'customer_id' => $customerId
+            'customer_id' => $customerId,
+            'default_billing' => $existingDefaultBilling,
+            'default_shipping' => $existingDefaultShipping,
         ];
         $newAddressData = [
             'country_id' => $countryId,
@@ -486,8 +492,8 @@ public function testExecute(
             ->method('getParam')
             ->willReturnMap([
                 ['id', null, $addressId],
-                ['default_billing', false, $addressId],
-                ['default_shipping', false, $addressId],
+                ['default_billing', $existingDefaultBilling, $setDefaultBilling],
+                ['default_shipping', $existingDefaultShipping, $setDefaultShipping],
             ]);
 
         $this->addressRepository->expects($this->once())
@@ -565,11 +571,11 @@ public function testExecute(
             ->willReturnSelf();
         $this->addressData->expects($this->once())
             ->method('setIsDefaultBilling')
-            ->with()
+            ->with($setDefaultBilling)
             ->willReturnSelf();
         $this->addressData->expects($this->once())
             ->method('setIsDefaultShipping')
-            ->with()
+            ->with($setDefaultShipping)
             ->willReturnSelf();
 
         $this->messageManager->expects($this->once())
@@ -628,11 +634,11 @@ public function dataProviderTestExecute(): array
 
             [1, 1, 1, 2, null, null, 12, null, null],
             [1, 1, 1, 2, 'Alaska', null, 12, null, 'CA'],
-            [1, 1, 1, 2, 'Alaska', 'AK', 12, 'California', null],
+            [1, 1, 1, 2, 'Alaska', 'AK', 12, 'California', null, true, true, true, false],
 
-            [1, 1, 1, 2, null, null, 12, null, null],
-            [1, 1, 1, 2, 'Alaska', null, 12, null, 'CA'],
-            [1, 1, 1, 2, 'Alaska', 'AK', 12, 'California', null],
+            [1, 1, 1, 2, null, null, 12, null, null, false, false, true, false],
+            [1, 1, 1, 2, 'Alaska', null, 12, null, 'CA', true, false, true, false],
+            [1, 1, 1, 2, 'Alaska', 'AK', 12, 'California', null, true, true, true, true],
         ];
     }