From 3a86bd1685939f759a3ee64535135469ab0b45b3 Mon Sep 17 00:00:00 2001
From: Tom Byers <tombaromba@gmail.com>
Date: Wed, 16 Feb 2022 15:56:52 +0000
Subject: [PATCH] Change internals of classesPersister

The assumption that the classes you want to
persist will always have parity with the elements
that have those classes, at that point, won't
always be true.

Because of that, this changes the way elements
with those classes are stored, to be in a map
between classes and the elements with them (at
that point).

Also includes an extra test for a scenario where
more than one updating component is in the page
with classes that need to persist through updates.
---
 app/assets/javascripts/updateContent.js | 30 ++++++++++-------
 tests/javascripts/updateContent.test.js | 44 ++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/app/assets/javascripts/updateContent.js b/app/assets/javascripts/updateContent.js
index 8931f285e..9ef216099 100644
--- a/app/assets/javascripts/updateContent.js
+++ b/app/assets/javascripts/updateContent.js
@@ -19,33 +19,39 @@
   // this can be removed in favour of a CSS-only solution.
   var classesPersister = {
     _classNames: [],
-    _$els: [],
+    _classesTo$ElsMap: {},
     addClassName: function (className) {
       if (this._classNames.indexOf(className) === -1) {
         this._classNames.push(className);
       }
     },
     remove: function () {
+      // Store references to any elements with class names to persist
       this._classNames.forEach(className => {
         var $elsWithClassName = $('.' + className).removeClass(className);
 
-        // store elements for that className at the same index
-        this._$els.push($elsWithClassName);
+        if ($elsWithClassName.length > 0) {
+          this._classesTo$ElsMap[className] = $elsWithClassName;
+        }
       });
     },
     replace: function () {
-      this._classNames.forEach((className, index) => {
-        var $el = this._$els[index];
+      var className;
 
-        // Avoid updating elements that are no longer present.
-        // elements removed will still exist in memory but won't be attached to the DOM any more
-        if (global.document.body.contains($el.get(0))) {
-          $el.addClass(className);
-        }
-      });
+      for (className in this._classesTo$ElsMap) {
+        this._classesTo$ElsMap[className].each((idx, el) => {
+
+          // Avoid updating elements that are no longer present.
+          // elements removed will still exist in memory but won't be attached to the DOM any more
+          if (global.document.body.contains(el)) {
+            $(el).addClass(className);
+          }
+
+        });
+      }
 
       // remove references to elements
-      this._$els = [];
+      this._classesTo$ElsMap = {};
     }
   };
 
diff --git a/tests/javascripts/updateContent.test.js b/tests/javascripts/updateContent.test.js
index 7db5ebacc..7c5492d7e 100644
--- a/tests/javascripts/updateContent.test.js
+++ b/tests/javascripts/updateContent.test.js
@@ -418,7 +418,7 @@ describe('Update content', () => {
 
     });
 
-    test("If other scripts have added classes to the DOM, they should persist through updates", () => {
+    test("If other scripts have added classes to the DOM, they should persist through updates to a single component", () => {
 
       document.body.innerHTML = getInitialHTMLString(getPartial(partialData));
 
@@ -451,6 +451,48 @@ describe('Update content', () => {
 
     });
 
+    test("If other scripts have added classes to the DOM, they should persist through updates to multiple components", () => {
+
+      // Create duplicate components in the page
+      document.body.innerHTML = getInitialHTMLString(getPartial(partialData)) + "\n" + getInitialHTMLString(getPartial(partialData));
+
+      var partialsInPage = document.querySelectorAll('.ajax-block-container');
+
+      // Mark classes to persist on the partials (2nd is made up)
+      partialsInPage[0].setAttribute('data-classes-to-persist', 'js-child-has-focus');
+      partialsInPage[1].setAttribute('data-classes-to-persist', 'js-2nd-child-has-focus');
+
+      // Add examples of those classes on each partial (2nd is made up)
+      partialsInPage[0].querySelectorAll('.file-list h2')[0].classList.add('js-child-has-focus');
+      partialsInPage[1].querySelectorAll('.file-list h2')[0].classList.add('js-2nd-child-has-focus');
+
+      // Add an item to trigger an update
+      partialData.push({
+        title: "Reservoir flooding template",
+        hint: "The local reservoir has flooded. All people within 5 miles should move to a safer location.",
+        status: "Waiting for approval",
+        areas: [
+          "Santa Claus Village, Rovaniemi A",
+          "Santa Claus Village, Rovaniemi D"
+        ]
+      });
+
+      // make all responses have an extra item
+      responseObj[updateKey] = getPartial(partialData);
+
+      // start the module
+      window.GOVUK.modules.start();
+      jest.advanceTimersByTime(2000);
+
+      // re-select in case nodes in partialsInPage have changed
+      partialsInPage = document.querySelectorAll('.ajax-block-container');
+
+      // check the classes are still there
+      expect(partialsInPage[0].querySelectorAll('.file-list h2')[0].classList.contains('js-child-has-focus')).toBe(true);
+      expect(partialsInPage[1].querySelectorAll('.file-list h2')[0].classList.contains('js-2nd-child-has-focus')).toBe(true);
+
+    });
+
   });
 
   afterEach(() => {