MAGETWO-52117: Customer group is not changed for logged in customer

Author: maxwhite-ua

app/code/Magento/Customer/Controller/Account/UpdateSession.php

@@ -70,15 +70,17 @@ public function __construct(
     public function execute()
     {
         $customerData = $this->jsonHelper->jsonDecode($this->getRequest()->getContent());
-        if (isset($customerData['customer_id']) && $this->notificationStorage->isExists(
+        if (isset($customerData['customer_id'])
+            && $this->notificationStorage->isExists(
                 NotificationStorage::UPDATE_CUSTOMER_SESSION,
-                $customerData['customer_id'])
+                $customerData['customer_id']
+            )
         ) {
             $customer = $this->customerRepository->getById($customerData['customer_id']);
             $this->session->setCustomerData($customer);
             $this->session->setCustomerGroupId($customer->getGroupId());
             $this->session->regenerateId();
-            $this->notificationStorage->remove(NotificationStorage::UPDATE_CUSTOMER_SESSION, $customerData['customer_id']);
+            $this->notificationStorage->remove(NotificationStorage::UPDATE_CUSTOMER_SESSION, $customer->getId());
         }
     }
 }

app/code/Magento/Customer/Model/Plugin/CustomerNotification.php

@@ -45,10 +45,13 @@ class CustomerNotification
     private $state;
 
     /**
+     * CustomerNotification constructor.
+     * 
      * @param Session $session
      * @param NotificationStorage $notificationStorage
      * @param CookieManagerInterface $cookieManager
      * @param CookieMetadataFactory $cookieMetadataFactory
+     * @param State $state
      */
     public function __construct(
         Session $session,