MAGETWO-80502: Product set to one website resets incorrectly for all websites after special price scheduled update ends

Author: dhorytskyi

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

@@ -292,27 +292,34 @@ protected function initializeProductData(array $productData, $createNew)
         foreach ($productData as $key => $value) {
             $product->setData($key, $value);
         }
-        $this->assignProductToWebsites($product);
+        $this->assignProductToWebsites($product, $createNew);
 
         return $product;
     }
 
     /**
      * @param \Magento\Catalog\Model\Product $product
+     * @param  bool $createNew
      * @return void
      */
-    private function assignProductToWebsites(\Magento\Catalog\Model\Product $product)
+    private function assignProductToWebsites(\Magento\Catalog\Model\Product $product, $createNew)
     {
-        if (!$this->storeManager->hasSingleStore()) {
-
-            if ($this->storeManager->getStore()->getCode() == \Magento\Store\Model\Store::ADMIN_CODE) {
-                $websiteIds = array_keys($this->storeManager->getWebsites());
-            } else {
-                $websiteIds = [$this->storeManager->getStore()->getWebsiteId()];
-            }
+        $websiteIds = $product->getWebsiteIds();
+
+        if ($createNew && !$this->storeManager->hasSingleStore()) {
+            $websiteIds = array_unique(
+                array_merge(
+                    $websiteIds,
+                    [$this->storeManager->getStore()->getWebsiteId()]
+                )
+            );
+        }
 
-            $product->setWebsiteIds(array_unique(array_merge($product->getWebsiteIds(), $websiteIds)));
+        if ($createNew && $this->storeManager->getStore(true)->getCode() == \Magento\Store\Model\Store::ADMIN_CODE) {
+            $websiteIds = array_keys($this->storeManager->getWebsites());
         }
+
+        $product->setWebsiteIds($websiteIds);
     }
 
     /**

app/code/Magento/Catalog/Test/Unit/Model/ProductRepositoryTest.php

@@ -12,6 +12,7 @@
 use Magento\Framework\Api\Data\ImageContentInterface;
 use Magento\Framework\Api\SortOrder;
 use Magento\Framework\TestFramework\Unit\Helper\ObjectManager;
+use Magento\Store\Api\Data\StoreInterface;
 
 /**
  * Tests \Magento\Catalog\Model\ProductRepositoryTest
@@ -1267,6 +1268,45 @@ public function testSaveExistingWithNewMediaGalleryEntries()
         $this->model->save($this->productMock);
     }
 
+    public function testSaveWithDifferentWebsites()
+    {
+        $getWebsitesResultData = [
+            1 => ['first'],
+            2 => ['second'],
+            3 => ['third']
+        ];
+        $getWebsiteIdsResultData = [1,2,3];
+        $setWebsiteIdsResultData = [2,3];
+        $getIdBySkuResultData = 100;
+        $storeMock = $this->getMock(StoreInterface::class);
+        $this->resourceModelMock->expects($this->at(0))->method('getIdBySku')->will($this->returnValue(null));
+        $this->resourceModelMock
+            ->expects($this->at(3))
+            ->method('getIdBySku')
+            ->will($this->returnValue($getIdBySkuResultData));
+        $this->productFactoryMock->expects($this->any())
+            ->method('create')
+            ->will($this->returnValue($this->productMock));
+        $this->initializationHelperMock->expects($this->never())->method('initialize');
+        $this->resourceModelMock->expects($this->once())->method('validate')->with($this->productMock)
+            ->willReturn(true);
+        $this->resourceModelMock->expects($this->once())->method('save')->with($this->productMock)->willReturn(true);
+        $this->extensibleDataObjectConverterMock
+            ->expects($this->once())
+            ->method('toNestedArray')
+            ->will($this->returnValue($this->productData));
+        $this->storeManagerMock->expects($this->any())
+            ->method('getStore')
+            ->willReturn($storeMock);
+        $this->storeManagerMock->expects($this->once())
+            ->method('getWebsites')
+            ->willReturn($getWebsitesResultData);
+        $this->productMock->expects($this->once())->method('getWebsiteIds')->willReturn($getWebsiteIdsResultData);
+        $this->productMock->expects($this->once())->method('setWebsiteIds')->willReturn($setWebsiteIdsResultData);
+
+        $this->assertEquals($this->productMock, $this->model->save($this->productMock));
+    }
+
     public function testSaveExistingWithMediaGalleryEntries()
     {
         //update one entry, delete one entry