Only load attributes relevant to the product's current attribute set

app/code/Magento/Catalog/Model/Entity/GetCategoryCustomAttributeCodes.php

@@ -29,7 +29,7 @@ public function __construct(
     /**
      * @inheritdoc
      */
-    public function execute(MetadataServiceInterface $metadataService): array
+    public function execute(MetadataServiceInterface $metadataService, ?int $attributeSetId = null): array
     {
         $customAttributesCodes = $this->baseCustomAttributeCodes->execute($metadataService);
         return array_diff($customAttributesCodes, CategoryInterface::ATTRIBUTES);

app/code/Magento/Catalog/Model/Entity/GetProductCustomAttributeCodes.php

@@ -7,8 +7,10 @@
 namespace Magento\Catalog\Model\Entity;
 
 use Magento\Catalog\Api\Data\ProductInterface;
+use Magento\Catalog\Api\ProductAttributeRepositoryInterface;
 use Magento\Eav\Model\Entity\GetCustomAttributeCodesInterface;
 use Magento\Framework\Api\MetadataServiceInterface;
+use Magento\Framework\Api\SearchCriteriaBuilder;
 
 class GetProductCustomAttributeCodes implements GetCustomAttributeCodesInterface
 {
@@ -18,20 +20,75 @@ class GetProductCustomAttributeCodes implements GetCustomAttributeCodesInterface
     private $baseCustomAttributeCodes;
 
     /**
-     * @param GetCustomAttributeCodesInterface $baseCustomAttributeCodes
+     * @var SearchCriteriaBuilder
+     */
+    private $searchCriteriaBuilder;
+
+    /**
+     * @var ProductAttributeRepositoryInterface
+     */
+    private $attributeRepository;
+
+    /**
+     * @var array[]
+     */
+    private $customAttributeCodes = [];
+
+    /**
+     * @param GetCustomAttributeCodesInterface    $baseCustomAttributeCodes
+     * @param SearchCriteriaBuilder               $searchCriteriaBuilder
+     * @param ProductAttributeRepositoryInterface $attributeRepository
      */
     public function __construct(
-        GetCustomAttributeCodesInterface $baseCustomAttributeCodes
+        GetCustomAttributeCodesInterface $baseCustomAttributeCodes,
+        SearchCriteriaBuilder $searchCriteriaBuilder,
+        ProductAttributeRepositoryInterface $attributeRepository
     ) {
         $this->baseCustomAttributeCodes = $baseCustomAttributeCodes;
+        $this->searchCriteriaBuilder = $searchCriteriaBuilder;
+        $this->attributeRepository = $attributeRepository;
     }
 
     /**
      * @inheritdoc
      */
-    public function execute(MetadataServiceInterface $metadataService): array
+    public function execute(MetadataServiceInterface $metadataService, ?int $attributeSetId = null): array
+    {
+        if (null !== $attributeSetId) {
+            return $this->getAttributesForSet($attributeSetId);
+        }
+
+        return $this->getAttributes($metadataService);
+    }
+
+    private function getAttributesForSet(int $attributeSetId)
+    {
+        if (!isset($this->customAttributeCodes[$attributeSetId])) {
+            $codes = [];
+            $criteria = $this->searchCriteriaBuilder->addFilter('attribute_set_id', $attributeSetId, 'eq');
+            $attributes = $this->attributeRepository->getList($criteria->create())->getItems();
+
+            if (is_array($attributes)) {
+                /** @var $attribute \Magento\Framework\Api\MetadataObjectInterface */
+                foreach ($attributes as $attribute) {
+                    $codes[] = $attribute->getAttributeCode();
+                }
+            }
+
+            $codes = array_values(
+                array_diff($codes, ProductInterface::ATTRIBUTES)
+            );
+
+            $this->customAttributeCodes[$attributeSetId] = $codes;
+        }
+
+        return $this->customAttributeCodes[$attributeSetId];
+    }
+
+    private function getAttributes(MetadataServiceInterface $metadataService)
     {
         $customAttributesCodes = $this->baseCustomAttributeCodes->execute($metadataService);
+
         return array_diff($customAttributesCodes, ProductInterface::ATTRIBUTES);
     }
 }

app/code/Magento/Catalog/Model/Product.php

@@ -497,7 +497,7 @@ protected function _getResource()
      */
     protected function getCustomAttributesCodes()
     {
-        return $this->getCustomAttributeCodes->execute($this->metadataService);
+        return $this->getCustomAttributeCodes->execute($this->metadataService, $this->getAttributeSetId());
     }
 
     /**

app/code/Magento/Catalog/Test/Unit/Model/Entity/GetProductCustomAttributeCodesTest.php

@@ -6,9 +6,14 @@
 
 namespace Magento\Catalog\Test\Unit\Model\Entity;
 
+use Magento\Catalog\Api\Data\ProductAttributeInterface;
+use Magento\Catalog\Api\Data\ProductAttributeSearchResultsInterface;
+use Magento\Catalog\Api\ProductAttributeRepositoryInterface;
 use Magento\Catalog\Model\Entity\GetProductCustomAttributeCodes;
 use Magento\Eav\Model\Entity\GetCustomAttributeCodesInterface;
 use Magento\Framework\Api\MetadataServiceInterface;
+use Magento\Framework\Api\SearchCriteria;
+use Magento\Framework\Api\SearchCriteriaBuilder;
 use Magento\Framework\TestFramework\Unit\Helper\ObjectManager;
 use PHPUnit\Framework\TestCase;
 
@@ -29,6 +34,16 @@ class GetProductCustomAttributeCodesTest extends TestCase
      */
     private $baseCustomAttributeCodes;
 
+    /**
+     * @var SearchCriteriaBuilder|\PHPUnit_Framework_MockObject_MockObject
+     */
+    private $searchCriteriaBuilder;
+
+    /**
+     * @var ProductAttributeRepositoryInterface|\PHPUnit_Framework_MockObject_MockObject
+     */
+    private $attributeRepository;
+
     /**
      * @inheritdoc
      */
@@ -38,10 +53,24 @@ protected function setUp()
             ->disableOriginalConstructor()
             ->setMethods(['execute'])
             ->getMockForAbstractClass();
+
+        $this->searchCriteriaBuilder = $this->getMockBuilder(SearchCriteriaBuilder::class)
+            ->disableOriginalConstructor()
+            ->getMock();
+
+        $this->attributeRepository = $this->getMockBuilder(ProductAttributeRepositoryInterface::class)
+            ->disableOriginalConstructor()
+            ->setMethods(['getList'])
+            ->getMockForAbstractClass();
+
         $objectManager = new ObjectManager($this);
         $this->getProductCustomAttributeCodes = $objectManager->getObject(
             GetProductCustomAttributeCodes::class,
-            ['baseCustomAttributeCodes' => $this->baseCustomAttributeCodes]
+            [
+                'baseCustomAttributeCodes' => $this->baseCustomAttributeCodes,
+                'searchCriteriaBuilder' => $this->searchCriteriaBuilder,
+                'attributeRepository' => $this->attributeRepository
+            ]
         );
     }
 
@@ -50,17 +79,156 @@ protected function setUp()
      */
     public function testExecute()
     {
-        /** @var MetadataServiceInterface|\PHPUnit_Framework_MockObject_MockObject $metadataService */
-        $metadataService = $this->getMockBuilder(MetadataServiceInterface::class)
-            ->disableOriginalConstructor()
-            ->getMockForAbstractClass();
+        $metadataService = $this->getMockMetadataService();
+
         $this->baseCustomAttributeCodes->expects($this->once())
             ->method('execute')
             ->with($this->identicalTo($metadataService))
             ->willReturn(['test_custom_attribute_code', 'name']);
+
         $this->assertEquals(
             ['test_custom_attribute_code'],
             $this->getProductCustomAttributeCodes->execute($metadataService)
         );
     }
+
+    public function testExecuteForAttributesInASet()
+    {
+        $metadataService = $this->getMockMetadataService();
+
+        $this->baseCustomAttributeCodes->expects($this->never())->method('execute');
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('addFilter')
+            ->with('attribute_set_id', 99, 'eq')
+            ->willReturnSelf();
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('create')
+            ->willReturn($this->getMockSearchCriteria());
+
+        $result = $this->getMockBuilder(ProductAttributeSearchResultsInterface::class)
+            ->setMethods(['getItems'])
+            ->getMockForAbstractClass();
+
+        $this->attributeRepository->expects($this->once())
+            ->method('getList')
+            ->willReturn($result);
+
+        $result->expects($this->once())
+            ->method('getItems')
+            ->willReturn(
+                [
+                    $this->getMockAttribute('test_code'),
+                    $this->getMockAttribute('another_code'),
+                    $this->getMockAttribute('sku'),
+                    $this->getMockAttribute('price')
+                ]
+            );
+
+        $this->assertEquals(
+            ['test_code', 'another_code'],
+            $this->getProductCustomAttributeCodes->execute($metadataService, 99)
+        );
+    }
+
+    public function testExecuteForAttributesInASetMemoizesResult()
+    {
+        $metadataService = $this->getMockMetadataService();
+
+        $this->baseCustomAttributeCodes->expects($this->never())->method('execute');
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('addFilter')
+            ->with('attribute_set_id', 100, 'eq')
+            ->willReturnSelf();
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('create')
+            ->willReturn($this->getMockSearchCriteria());
+
+        $result = $this->getMockBuilder(ProductAttributeSearchResultsInterface::class)
+            ->setMethods(['getItems'])
+            ->getMockForAbstractClass();
+
+        $this->attributeRepository->expects($this->once())
+            ->method('getList')
+            ->willReturn($result);
+
+        $result->expects($this->once())
+            ->method('getItems')
+            ->willReturn([$this->getMockAttribute('foo'), $this->getMockAttribute('price')]);
+
+        $this->assertEquals(
+            ['foo'],
+            $this->getProductCustomAttributeCodes->execute($metadataService, 100)
+        );
+
+        $this->assertEquals(
+            ['foo'],
+            $this->getProductCustomAttributeCodes->execute($metadataService, 100)
+        );
+    }
+
+    public function testExecuteForAttributesInASetReturnsEmptyArrayWhenNoAttributesFound()
+    {
+        $metadataService = $this->getMockMetadataService();
+
+        $this->baseCustomAttributeCodes->expects($this->never())->method('execute');
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('addFilter')
+            ->with('attribute_set_id', 101, 'eq')
+            ->willReturnSelf();
+
+        $this->searchCriteriaBuilder->expects($this->once())
+            ->method('create')
+            ->willReturn($this->getMockSearchCriteria());
+
+        $result = $this->getMockBuilder(ProductAttributeSearchResultsInterface::class)
+            ->setMethods(['getItems'])
+            ->getMockForAbstractClass();
+
+        $this->attributeRepository->expects($this->once())
+            ->method('getList')
+            ->willReturn($result);
+
+        $result->expects($this->once())
+            ->method('getItems')
+            ->willReturn([]);
+
+        $this->assertEquals(
+            [],
+            $this->getProductCustomAttributeCodes->execute($metadataService, 101)
+        );
+    }
+
+    /**
+     * @return MetadataServiceInterface|\PHPUnit_Framework_MockObject_MockObject
+     */
+    private function getMockMetadataService()
+    {
+        return $this->getMockBuilder(MetadataServiceInterface::class)
+            ->disableOriginalConstructor()
+            ->getMockForAbstractClass();
+    }
+
+    private function getMockSearchCriteria()
+    {
+        return $this->getMockBuilder(SearchCriteria::class)->disableOriginalConstructor()->getMock();
+    }
+
+    private function getMockAttribute($code)
+    {
+        $attribute = $this->getMockBuilder(ProductAttributeInterface::class)
+            ->disableOriginalConstructor()
+            ->setMethods(['getAttributeCode'])
+            ->getMockForAbstractClass();
+
+        $attribute->expects($this->any())
+            ->method('getAttributeCode')
+            ->willReturn($code);
+
+        return $attribute;
+    }
 }

app/code/Magento/Eav/Model/Entity/GetCustomAttributeCodes.php

@@ -19,9 +19,12 @@ class GetCustomAttributeCodes implements GetCustomAttributeCodesInterface
      * Receive a list of custom EAV attributes using provided metadata service. The results are cached per entity type
      *
      * @param MetadataServiceInterface $metadataService Custom attribute metadata service to be used
+     * @param int|null                 $attributeSetId  Optional attribute set ID, if provided will only load attributes
+     *                                                  for that attribute set.
+     *
      * @return string[]
      */
-    public function execute(MetadataServiceInterface $metadataService): array
+    public function execute(MetadataServiceInterface $metadataService, ?int $attributeSetId = null): array
     {
         $cacheKey = get_class($metadataService);
         if (!isset($this->customAttributesCodes[$cacheKey])) {

app/code/Magento/Eav/Model/Entity/GetCustomAttributeCodesInterface.php

@@ -14,7 +14,9 @@ interface GetCustomAttributeCodesInterface
      * Receive a list of custom EAV attributes using provided metadata service.
      *
      * @param MetadataServiceInterface $metadataService Custom attribute metadata service to be used
+     * @param int|null                 $attributeSetId
+     *
      * @return string[]
      */
-    public function execute(MetadataServiceInterface $metadataService): array;
+    public function execute(MetadataServiceInterface $metadataService, ?int $attributeSetId = null): array;
 }