ACP2E-2213: Loading order through repository loads the order multiple times from DB

Author: o-iegorov

app/code/Magento/Sales/Model/OrderRepository.php

@@ -93,7 +93,8 @@ class OrderRepository implements \Magento\Sales\Api\OrderRepositoryInterface, Re
      * @param OrderTaxManagementInterface|null $orderTaxManagement
      * @param PaymentAdditionalInfoInterfaceFactory|null $paymentAdditionalInfoFactory
      * @param JsonSerializer|null $serializer
-     * @param JoinProcessorInterface $extensionAttributesJoinProcessor
+     * @param JoinProcessorInterface|null $extensionAttributesJoinProcessor
+     * @param ShippingAssignmentBuilder|null $shippingAssignmentBuilder
      */
     public function __construct(
         Metadata $metadata,
@@ -103,7 +104,8 @@ public function __construct(
         OrderTaxManagementInterface $orderTaxManagement = null,
         PaymentAdditionalInfoInterfaceFactory $paymentAdditionalInfoFactory = null,
         JsonSerializer $serializer = null,
-        JoinProcessorInterface $extensionAttributesJoinProcessor = null
+        JoinProcessorInterface $extensionAttributesJoinProcessor = null,
+        ShippingAssignmentBuilder $shippingAssignmentBuilder = null
     ) {
         $this->metadata = $metadata;
         $this->searchResultFactory = $searchResultFactory;
@@ -119,6 +121,8 @@ public function __construct(
             ->get(JsonSerializer::class);
         $this->extensionAttributesJoinProcessor = $extensionAttributesJoinProcessor
             ?: ObjectManager::getInstance()->get(JoinProcessorInterface::class);
+        $this->shippingAssignmentBuilder = $shippingAssignmentBuilder
+            ?: ObjectManager::getInstance()->get(ShippingAssignmentBuilder::class);
     }
 
     /**

app/code/Magento/Sales/Test/Unit/Model/OrderRepositoryTest.php

@@ -22,6 +22,7 @@
 use Magento\Sales\Model\Order\Shipping;
 use Magento\Sales\Model\Order\ShippingAssignment;
 use Magento\Sales\Model\Order\ShippingAssignmentBuilder;
+use Magento\Sales\Api\Data\OrderAddressInterface;
 use Magento\Sales\Model\OrderRepository;
 use Magento\Sales\Model\ResourceModel\Metadata;
 use Magento\Sales\Model\ResourceModel\Order as OrderResource;
@@ -77,6 +78,11 @@ class OrderRepositoryTest extends TestCase
      */
     private $orderExtensionFactoryMock;
 
+    /**
+     * @var ShippingAssignmentBuilder|MockObject
+     */
+    private $shippingAssignmentBuilder;
+
     /**
      * Setup the test
      *
@@ -104,6 +110,9 @@ protected function setUp(): void
         $this->paymentAdditionalInfoFactory = $this->getMockBuilder(PaymentAdditionalInfoInterfaceFactory::class)
             ->disableOriginalConstructor()
             ->setMethods(['create'])->getMockForAbstractClass();
+        $this->shippingAssignmentBuilder = $this->getMockBuilder(ShippingAssignmentBuilder::class)
+            ->disableOriginalConstructor()
+            ->getMock();
         $this->orderRepository = $this->objectManager->getObject(
             OrderRepository::class,
             [
@@ -113,6 +122,7 @@ protected function setUp(): void
                 'orderExtensionFactory' => $this->orderExtensionFactoryMock,
                 'orderTaxManagement' => $this->orderTaxManagementMock,
                 'paymentAdditionalInfoFactory' => $this->paymentAdditionalInfoFactory,
+                'shippingAssignmentBuilder' => $this->shippingAssignmentBuilder
             ]
         );
     }
@@ -194,7 +204,10 @@ public function testSave()
             ->method('getShippingAssignments')
             ->willReturn([$shippingAssignment]);
         $shippingAssignment->expects($this->once())->method('getShipping')->willReturn($shippingMock);
-        $shippingMock->expects($this->once())->method('getAddress');
+        $shippingAddressMock = $this->getMockBuilder(OrderAddressInterface::class)
+            ->disableOriginalConstructor()
+            ->getMock();
+        $shippingMock->expects($this->once())->method('getAddress')->willReturn($shippingAddressMock);
         $shippingMock->expects($this->once())->method('getMethod');
         $this->metadata->expects($this->once())->method('getMapper')->willReturn($mapperMock);
         $mapperMock->expects($this->once())->method('save');
@@ -219,7 +232,7 @@ public function testGet()
             ->disableOriginalConstructor()->getMockForAbstractClass();
         $paymentMock->expects($this->once())->method('getAdditionalInformation')->willReturn($paymentInfo);
         $orderExtension = $this->getOrderExtensionMock();
-        $orderExtension->expects($this->once())->method('getShippingAssignments')->willReturn(true);
+        $orderExtension->expects($this->once())->method('getShippingAssignments')->willReturn(null);
         $orderExtension->expects($this->once())->method('setAppliedTaxes')->with($appliedTaxes);
         $orderExtension->expects($this->once())->method('setConvertingFromQuote')->with(true);
         $orderExtension->expects($this->once())->method('setItemAppliedTaxes')->with($items);
@@ -228,7 +241,7 @@ public function testGet()
         $orderEntity->expects($this->once())->method('load')->with($orderId)->willReturn($orderEntity);
         $orderEntity->expects($this->exactly(2))->method('getEntityId')->willReturn($orderId);
         $orderEntity->expects($this->once())->method('getPayment')->willReturn($paymentMock);
-        $orderEntity->expects($this->exactly(2))->method('setExtensionAttributes')->with($orderExtension);
+        $orderEntity->expects($this->exactly(3))->method('setExtensionAttributes')->with($orderExtension);
         $orderEntity->expects($this->exactly(3))
             ->method('getExtensionAttributes')
             ->willReturnOnConsecutiveCalls(null, $orderExtension, $orderExtension);
@@ -240,6 +253,7 @@ public function testGet()
         $orderTaxDetailsMock->expects($this->once())->method('getItems')->willReturn($items);
         $this->orderTaxManagementMock->expects($this->atLeastOnce())->method('getOrderTaxDetails')
             ->willReturn($orderTaxDetailsMock);
+        $this->shippingAssignmentBuilder->expects($this->once())->method('setOrder')->with($orderEntity);
 
         $this->orderRepository->get($orderId);
     }
@@ -254,7 +268,7 @@ private function getOrderExtensionMock(): MockObject
         $mockBuilder = $this->getMockBuilder(OrderExtensionInterface::class)->disableOriginalConstructor();
         try {
             $mockBuilder
-                ->addMethods(
+                ->onlyMethods(
                     [
                         'getShippingAssignments',
                         'setAppliedTaxes',