ACP2E-3433: Allow Countries configuration causes issues in customer address configurations

- fix
- add test

Author: o-dubovyk

app/code/Magento/Customer/Model/ResourceModel/Address/Attribute/Source/Country.php

@@ -1,12 +1,14 @@
 <?php
 /**
- * Copyright © Magento, Inc. All rights reserved.
- * See COPYING.txt for license details.
+ * Copyright 2013 Adobe
+ * All Rights Reserved.
  */
 
 namespace Magento\Customer\Model\ResourceModel\Address\Attribute\Source;
 
+use Magento\Framework\App\Area;
 use Magento\Framework\App\ObjectManager;
+use Magento\Store\Model\Store;
 use Magento\Store\Model\StoreManagerInterface;
 
 /**
@@ -24,17 +26,25 @@ class Country extends \Magento\Eav\Model\Entity\Attribute\Source\Table
      */
     private $storeManager;
 
+    /**
+     * @var \Magento\Framework\App\State
+     */
+    private $state;
+
     /**
      * @param \Magento\Eav\Model\ResourceModel\Entity\Attribute\Option\CollectionFactory $attrOptionCollectionFactory
      * @param \Magento\Eav\Model\ResourceModel\Entity\Attribute\OptionFactory $attrOptionFactory
      * @param \Magento\Directory\Model\ResourceModel\Country\CollectionFactory $countriesFactory
+     * @param \Magento\Framework\App\State $state
      */
     public function __construct(
         \Magento\Eav\Model\ResourceModel\Entity\Attribute\Option\CollectionFactory $attrOptionCollectionFactory,
         \Magento\Eav\Model\ResourceModel\Entity\Attribute\OptionFactory $attrOptionFactory,
-        \Magento\Directory\Model\ResourceModel\Country\CollectionFactory $countriesFactory
+        \Magento\Directory\Model\ResourceModel\Country\CollectionFactory $countriesFactory,
+        \Magento\Framework\App\State $state
     ) {
         $this->_countriesFactory = $countriesFactory;
+        $this->state = $state;
         parent::__construct($attrOptionCollectionFactory, $attrOptionFactory);
     }
 
@@ -44,9 +54,14 @@ public function __construct(
     public function getAllOptions($withEmpty = true, $defaultValues = false)
     {
         if (!$this->_options) {
-            $this->_options = $this->_createCountriesCollection()->loadByStore(
-                $this->getStoreManager()->getStore()->getId()
-            )->toOptionArray();
+            if ($this->state->getAreaCode() === Area::AREA_ADMINHTML) {
+                $storeId = $this->getAttribute()->getWebsite() &&
+                    $this->getAttribute()->getWebsite()->getDefaultStore() ?
+                    $this->getAttribute()->getWebsite()->getDefaultStore()->getId() : Store::DEFAULT_STORE_ID;
+            } else {
+                $storeId = $this->getStoreManager()->getStore()->getId();
+            }
+            $this->_options = $this->_createCountriesCollection()->loadByStore($storeId)->toOptionArray();
         }
         return $this->_options;
     }

dev/tests/integration/testsuite/Magento/Customer/Model/ResourceModel/Address/Attribute/Source/CountryTest.php

@@ -0,0 +1,88 @@
+<?php
+/**
+ * Copyright 2024 Adobe
+ * All Rights Reserved.
+ */
+declare(strict_types=1);
+
+namespace Magento\Customer\Model\ResourceModel\Address\Attribute\Source;
+
+use Magento\Eav\Model\Config;
+use Magento\Framework\ObjectManagerInterface;
+use Magento\Framework\Validator\UniversalFactory;
+use Magento\Store\Model\StoreManagerInterface;
+use Magento\TestFramework\Helper\Bootstrap;
+
+class CountryTest extends \PHPUnit\Framework\TestCase
+{
+    /**
+     * @var ObjectManagerInterface
+     */
+    private $objectManager;
+
+    /**
+     * @var StoreManagerInterface
+     */
+    private $storeManager;
+
+    /**
+     * @var Config
+     */
+    private $eavConfig;
+
+    /**
+     * @var UniversalFactory
+     */
+    private $factory;
+
+    /**
+     * {@inheritdoc}
+     */
+    public function setUp(): void
+    {
+        $this->objectManager = Bootstrap::getObjectManager();
+        $this->storeManager = $this->objectManager->get(StoreManagerInterface::class);
+        $this->eavConfig = $this->objectManager->get(Config::class);
+        $this->factory = $this->objectManager->get(UniversalFactory::class);
+
+        parent::setUp();
+    }
+
+    /**
+     * Assert that countries are returned according to allowed countries per respective website
+     *
+     * @magentoAppArea adminhtml
+     * @magentoDbIsolation disabled
+     * @magentoDataFixture Magento/Store/_files/second_website_with_store_group_and_store.php
+     * @magentoConfigFixture default_store general/country/default NL
+     * @magentoConfigFixture default_store general/country/allow NL
+     * @magentoConfigFixture fixture_second_store_store general/country/default UA
+     * @magentoConfigFixture fixture_second_store_store general/country/allow UA
+     * @dataProvider dataProvider
+     */
+    public function testMethod($website, $expectedCountryCont, $expectedCountryCode)
+    {
+        $websiteId = $this->storeManager->getWebsite($website)->getId();
+        $attribute = $this->eavConfig->getAttribute('customer_address', 'country_id');
+        $attribute->setWebsite($websiteId);
+
+        $countryOptions = $this->factory->create($attribute->getSourceModel())
+            ->setAttribute($attribute)->getAllOptions();
+
+        $countryOptions = array_map(fn($countryOption) => $countryOption['value'], $countryOptions);
+        $this->assertEquals($expectedCountryCont, count($countryOptions));
+        $this->assertContains($expectedCountryCode, $countryOptions);
+    }
+
+    /**
+     * @return array
+     */
+    public static function dataProvider()
+    {
+        return [
+            ['admin', 249, 'GB'],
+            ['base', 1, 'NL'],
+            ['test', 1, 'UA'],
+        ];
+    }
+}