ENGCOM-3082: Fix Customer custom attributes lost after save #17968

 - Merge Pull Request #17968 from Thundar/magento2:issue/12479
 - Merged commits:
   1. 7965c9b
   2. 44fe03f
   3. 6a098a7

Author: magento-engcom-team

app/code/Magento/Customer/Model/Address.php

@@ -154,9 +154,6 @@ public function updateData(AddressInterface $address)
         // Need to explicitly set this due to discrepancy in the keys between model and data object
         $this->setIsDefaultBilling($address->isDefaultBilling());
         $this->setIsDefaultShipping($address->isDefaultShipping());
-        if (!$this->getAttributeSetId()) {
-            $this->setAttributeSetId(AddressMetadataInterface::ATTRIBUTE_SET_ID_ADDRESS);
-        }
         $customAttributes = $address->getCustomAttributes();
         if ($customAttributes !== null) {
             foreach ($customAttributes as $attribute) {

app/code/Magento/Customer/Model/Customer.php

@@ -358,13 +358,6 @@ public function updateData($customer)
             $this->setId($customerId);
         }
 
-        // Need to use attribute set or future updates can cause data loss
-        if (!$this->getAttributeSetId()) {
-            $this->setAttributeSetId(
-                CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER
-            );
-        }
-
         return $this;
     }
 
@@ -960,6 +953,16 @@ public function getSharedWebsiteIds()
         return $ids;
     }
 
+    /**
+     * Retrieve attribute set id for customer.
+     *
+     * @return int
+     */
+    public function getAttributeSetId()
+    {
+        return parent::getAttributeSetId() ?: CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER;
+    }
+
     /**
      * Set store to customer
      *

app/code/Magento/Customer/Model/ResourceModel/CustomerRepository.php

@@ -205,12 +205,6 @@ public function save(\Magento\Customer\Api\Data\CustomerInterface $customer, $pa
             $customerModel->setStoreId($this->storeManager->getStore()->getId());
         }
 
-        // Need to use attribute set or future updates can cause data loss
-        if (!$customerModel->getAttributeSetId()) {
-            $customerModel->setAttributeSetId(
-                \Magento\Customer\Api\CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER
-            );
-        }
         $this->populateCustomerWithSecureData($customerModel, $passwordHash);
 
         // If customer email was changed, reset RpToken info

app/code/Magento/Customer/Test/Unit/Model/CustomerTest.php

@@ -307,8 +307,6 @@ public function testUpdateData()
         }
 
         $expectedResult[$attribute->getAttributeCode()] = $attribute->getValue();
-        $expectedResult['attribute_set_id'] =
-            \Magento\Customer\Api\CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER;
 
         $this->assertEquals($this->_model->getData(), $expectedResult);
     }

app/code/Magento/Customer/Test/Unit/Model/ResourceModel/CustomerRepositoryTest.php

@@ -334,12 +334,6 @@ public function testSave()
         $customerModel->expects($this->once())
             ->method('setId')
             ->with($customerId);
-        $customerModel->expects($this->once())
-            ->method('getAttributeSetId')
-            ->willReturn(null);
-        $customerModel->expects($this->once())
-            ->method('setAttributeSetId')
-            ->with(\Magento\Customer\Api\CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER);
         $customerAttributesMetaData->expects($this->atLeastOnce())
             ->method('getId')
             ->willReturn($customerId);
@@ -616,12 +610,6 @@ public function testSaveWithPasswordHash()
         $customerModel->expects($this->once())
             ->method('setId')
             ->with($customerId);
-        $customerModel->expects($this->once())
-            ->method('getAttributeSetId')
-            ->willReturn(null);
-        $customerModel->expects($this->once())
-            ->method('setAttributeSetId')
-            ->with(\Magento\Customer\Api\CustomerMetadataInterface::ATTRIBUTE_SET_ID_CUSTOMER);
         $customerModel->expects($this->atLeastOnce())
             ->method('getId')
             ->willReturn($customerId);