Fix dark and white theme switch in Chrome

src/theme.ts

@@ -237,13 +237,13 @@ export async function setTheme(theme?: string): Promise<void> {
 
     // look for the stylesheet elements.
     // styleElements is a map from style name to HTMLLinkElement.
-    const styleElements = Object.create(null);
-    const themes = Array.from(document.querySelectorAll('[data-mx-theme]'));
+    const styleElements = new Map<string, HTMLLinkElement>();
+    const themes = Array.from(document.querySelectorAll<HTMLLinkElement>('[data-mx-theme]'));
     themes.forEach(theme => {
-        styleElements[theme.attributes['data-mx-theme'].value.toLowerCase()] = theme;
+        styleElements.set(theme.attributes['data-mx-theme'].value.toLowerCase(), theme);
     });
 
-    if (!(stylesheetName in styleElements)) {
+    if (!styleElements.has(stylesheetName)) {
         throw new Error("Unknown theme " + stylesheetName);
     }
 
@@ -258,17 +258,18 @@ export async function setTheme(theme?: string): Promise<void> {
     // having them interact badly... but this causes a flash of unstyled app
     // which is even uglier.  So we don't.
 
-    styleElements[stylesheetName].disabled = false;
+    const styleSheet = styleElements.get(stylesheetName);
+    styleSheet.disabled = false;
 
-    return new Promise((resolve) => {
+    return new Promise(((resolve, reject) => {
         const switchTheme = function() {
             // we re-enable our theme here just in case we raced with another
             // theme set request as per https://github.com/vector-im/element-web/issues/5601.
             // We could alternatively lock or similar to stop the race, but
             // this is probably good enough for now.
-            styleElements[stylesheetName].disabled = false;
-            Object.values(styleElements).forEach((a: HTMLStyleElement) => {
-                if (a == styleElements[stylesheetName]) return;
+            styleSheet.disabled = false;
+            styleElements.forEach(a => {
+                if (a == styleSheet) return;
                 a.disabled = true;
             });
             const bodyStyles = global.getComputedStyle(document.body);
@@ -279,26 +280,50 @@ export async function setTheme(theme?: string): Promise<void> {
             resolve();
         };
 
-        // turns out that Firefox preloads the CSS for link elements with
-        // the disabled attribute, but Chrome doesn't.
+        const isStyleSheetLoaded = () => Boolean(
+            [...document.styleSheets]
+                .find(_styleSheet => _styleSheet?.href === styleSheet.href),
+        );
 
-        let cssLoaded = false;
+        function waitForStyleSheetLoading() {
+            // turns out that Firefox preloads the CSS for link elements with
+            // the disabled attribute, but Chrome doesn't.
+            if (isStyleSheetLoaded()) {
+                switchTheme();
+                return;
+            }
 
-        styleElements[stylesheetName].onload = () => {
-            switchTheme();
-        };
+            let counter = 0;
 
-        for (let i = 0; i < document.styleSheets.length; i++) {
-            const ss = document.styleSheets[i];
-            if (ss && ss.href === styleElements[stylesheetName].href) {
-                cssLoaded = true;
-                break;
-            }
-        }
+            // In case of theme toggling (white => black => white)
+            // Chrome doesn't fire the `load` event when the white theme is selected the second times
+            const intervalId = setInterval(() => {
+                if (isStyleSheetLoaded()) {
+                    clearInterval(intervalId);
+                    styleSheet.onload = undefined;
+                    styleSheet.onerror = undefined;
+                    switchTheme();
+                }
 
-        if (cssLoaded) {
-            styleElements[stylesheetName].onload = undefined;
-            switchTheme();
+                // Avoid to be stuck in an endless loop if there is an issue in the stylesheet loading
+                counter++;
+                if (counter === 10) {
+                    clearInterval(intervalId);
+                    reject();
+                }
+            }, 200);
+
+            styleSheet.onload = () => {
+                clearInterval(intervalId);
+                switchTheme();
+            };
+
+            styleSheet.onerror = (e) => {
+                clearInterval(intervalId);
+                reject(e);
+            };
         }
-    });
+
+        waitForStyleSheetLoading();
+    }));
 }

test/theme-test.ts

@@ -0,0 +1,115 @@
+/*
+Copyright 2022 The Matrix.org Foundation C.I.C.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+import { setTheme } from "../src/theme";
+
+describe('theme', () => {
+    describe('setTheme', () => {
+        let lightTheme;
+        let darkTheme;
+
+        let spyQuerySelectorAll: jest.MockInstance<NodeListOf<Element>, [selectors: string]>;
+
+        beforeEach(() => {
+            const styles = [
+                {
+                    attributes: {
+                        'data-mx-theme': {
+                            value: 'light',
+                        },
+                    },
+                    disabled: true,
+                    href: 'urlLight',
+                    onload: () => void 0,
+                },
+                {
+                    attributes: {
+                        'data-mx-theme': {
+                            value: 'dark',
+                        },
+                    },
+                    disabled: true,
+                    href: 'urlDark',
+                    onload: () => void 0,
+                },
+            ];
+            lightTheme = styles[0];
+            darkTheme = styles[1];
+
+            jest.spyOn(document.body, 'style', 'get').mockReturnValue([] as any);
+            spyQuerySelectorAll = jest.spyOn(document, 'querySelectorAll').mockReturnValue(styles as any);
+        });
+
+        afterEach(() => {
+            jest.restoreAllMocks();
+            jest.useRealTimers();
+        });
+
+        it('should switch theme on onload call', async () => {
+            // When
+            await new Promise(resolve => {
+                setTheme('light').then(resolve);
+                lightTheme.onload();
+            });
+
+            // Then
+            expect(spyQuerySelectorAll).toHaveBeenCalledWith('[data-mx-theme]');
+            expect(spyQuerySelectorAll).toBeCalledTimes(1);
+            expect(lightTheme.disabled).toBe(false);
+            expect(darkTheme.disabled).toBe(true);
+        });
+
+        it('should reject promise on onerror call', () => {
+            return expect(new Promise(resolve => {
+                setTheme('light').catch(e => resolve(e));
+                lightTheme.onerror('call onerror');
+            })).resolves.toBe('call onerror');
+        });
+
+        it('should switch theme if CSS are preloaded', async () => {
+            // When
+            jest.spyOn(document, 'styleSheets', 'get').mockReturnValue([lightTheme] as any);
+
+            await setTheme('light');
+
+            // Then
+            expect(lightTheme.disabled).toBe(false);
+            expect(darkTheme.disabled).toBe(true);
+        });
+
+        it('should switch theme if CSS is loaded during pooling', async () => {
+            // When
+            jest.useFakeTimers();
+            await new Promise(resolve => {
+                setTheme('light').then(resolve);
+                jest.spyOn(document, 'styleSheets', 'get').mockReturnValue([lightTheme] as any);
+                jest.advanceTimersByTime(200);
+            });
+
+            // Then
+            expect(lightTheme.disabled).toBe(false);
+            expect(darkTheme.disabled).toBe(true);
+        });
+
+        it('should reject promise if pooling maximum value is reached', () => {
+            jest.useFakeTimers();
+            return new Promise(resolve => {
+                setTheme('light').catch(resolve);
+                jest.advanceTimersByTime(200 * 10);
+            });
+        });
+    });
+});