ACP2E-3382: Exported customer address cannot be imported

Author: abukatar

app/code/Magento/CustomerImportExport/Model/Import/Address.php

@@ -6,7 +6,7 @@
 
 namespace Magento\CustomerImportExport\Model\Import;
 
-use Magento\Customer\Model\ResourceModel\Address\Attribute\Source\CountryWithWebsites as CountryWithWebsitesSource;
+use Magento\CustomerImportExport\Model\Import\CountryWithWebsites as CountryWithWebsitesSource;
 use Magento\Eav\Model\Entity\Attribute\AbstractAttribute;
 use Magento\Framework\App\ObjectManager;
 use Magento\ImportExport\Model\Import\ErrorProcessing\ProcessingErrorAggregatorInterface;
@@ -365,7 +365,7 @@ public function getAttributeOptions(AbstractAttribute $attribute, array $indexAt
         if ($attribute->getAttributeCode() === 'country_id') {
             //If we want to get available options for country field then we have to use alternative source
             // to get actual data for each website.
-            $options = $this->countryWithWebsites->getAllOptions();
+            $options = $this->countryWithWebsites->getCountiesPerWebsite();
             //Available country options now will be sorted by websites.
             $code = $attribute->getAttributeCode();
             $websiteOptions = [Store::DEFAULT_STORE_ID => $standardOptions];

app/code/Magento/CustomerImportExport/Model/Import/CountryWithWebsites.php

@@ -0,0 +1,109 @@
+<?php
+/**
+ * Copyright 2024 Adobe
+ * All Rights Reserved.
+ */
+
+declare(strict_types=1);
+
+namespace Magento\CustomerImportExport\Model\Import;
+
+use Magento\Customer\Model\Config\Share as CustomerShareConfig;
+use Magento\Directory\Model\AllowedCountries;
+use Magento\Directory\Model\ResourceModel\Country\Collection as CountryCollection;
+use Magento\Directory\Model\ResourceModel\Country\CollectionFactory as CountryCollectionFactory;
+use Magento\Store\Model\ScopeInterface;
+use Magento\Store\Model\StoreManagerInterface;
+
+/**
+ * Return allowed countries for specified website for customer import
+ */
+class CountryWithWebsites
+{
+    /**
+     * @var CountryCollectionFactory
+     */
+    private $countriesFactory;
+
+    /**
+     * @var AllowedCountries
+     */
+    private $allowedCountriesReader;
+
+    /**
+     * @var array
+     */
+    private $allowedCountries;
+
+    /**
+     * @var StoreManagerInterface
+     */
+    private $storeManager;
+
+    /**
+     * @var CustomerShareConfig
+     */
+    private $shareConfig;
+
+    /**
+     * @param CountryCollectionFactory $countriesFactory
+     * @param AllowedCountries $allowedCountriesReader
+     * @param StoreManagerInterface $storeManager
+     * @param CustomerShareConfig $shareConfig
+     */
+    public function __construct(
+        CountryCollectionFactory $countriesFactory,
+        AllowedCountries $allowedCountriesReader,
+        StoreManagerInterface $storeManager,
+        CustomerShareConfig $shareConfig
+    ) {
+        $this->countriesFactory = $countriesFactory;
+        $this->allowedCountriesReader = $allowedCountriesReader;
+        $this->storeManager = $storeManager;
+        $this->shareConfig = $shareConfig;
+    }
+
+    /**
+     * Get allowed countries for specified website
+     *
+     * @return array
+     */
+    public function getCountiesPerWebsite(): array
+    {
+        if (!$this->allowedCountries) {
+            $websiteIds = [];
+
+            foreach ($this->storeManager->getWebsites() as $website) {
+                $countries = $this->allowedCountriesReader
+                    ->getAllowedCountries(ScopeInterface::SCOPE_WEBSITE, $website->getId());
+                $allowedCountries[] = $countries;
+
+                foreach ($countries as $countryCode) {
+                    $websiteIds[$countryCode][] = $website->getId();
+                }
+            }
+
+            $this->allowedCountries = $this->createCountriesCollection()
+                ->addFieldToFilter('country_id', ['in' => $allowedCountries])
+                ->toOptionArray();
+
+            foreach ($this->allowedCountries as &$option) {
+                if (isset($websiteIds[$option['value']])) {
+                    $option['website_ids'] = $websiteIds[$option['value']];
+                }
+            }
+        }
+
+        return $this->allowedCountries;
+    }
+
+    /**
+     * Create Countries Collection with all countries
+     *
+     * @return CountryCollection
+     */
+    private function createCountriesCollection()
+    {
+        return $this->countriesFactory->create();
+    }
+}

app/code/Magento/CustomerImportExport/Test/Unit/Model/Import/AddressTest.php

@@ -12,7 +12,7 @@
 use Magento\Customer\Model\CustomerFactory;
 use Magento\Customer\Model\Indexer\Processor;
 use Magento\Customer\Model\ResourceModel\Address\Attribute as AddressAttribute;
-use Magento\Customer\Model\ResourceModel\Address\Attribute\Source\CountryWithWebsites;
+use Magento\CustomerImportExport\Model\Import\CountryWithWebsites;
 use Magento\CustomerImportExport\Model\Import\Address;
 use Magento\CustomerImportExport\Model\ResourceModel\Import\Customer\Storage;
 use Magento\CustomerImportExport\Model\ResourceModel\Import\Customer\StorageFactory;
@@ -145,7 +145,7 @@ class AddressTest extends TestCase
     protected $errorAggregator;
 
     /**
-     * @var AddressAttribute\Source\CountryWithWebsites|MockObject
+     * @var CountryWithWebsites|MockObject
      */
     private $countryWithWebsites;
 
@@ -168,8 +168,7 @@ protected function setUp(): void
             ->disableOriginalConstructor()
             ->getMock();
         $this->countryWithWebsites
-
-            ->method('getAllOptions')
+            ->method('getCountiesPerWebsite')
             ->willReturn([]);
         $this->_model = $this->_getModelMock();
         $this->errorAggregator = $this->createPartialMock(