Add localStorage fallback class

src/common/auth/token_storage.ts

@@ -1,5 +1,8 @@
 import type { AuthData } from "home-assistant-js-websocket";
 import { extractSearchParam } from "../url/search-params";
+import { FallbackStorage } from "../../../test/test_helper/local-storage-fallback";
+
+const storage = window.localStorage || new FallbackStorage();
 
 declare global {
   interface Window {
@@ -36,7 +39,7 @@ export function saveTokens(tokens: AuthData | null) {
 
   if (tokenCache.writeEnabled) {
     try {
-      window.localStorage.setItem("hassTokens", JSON.stringify(tokens));
+      storage.setItem("hassTokens", JSON.stringify(tokens));
     } catch (err: any) {
       // write failed, ignore it. Happens if storage is full or private mode.
       // eslint-disable-next-line no-console
@@ -59,7 +62,7 @@ export function enableWrite() {
 export function loadTokens() {
   if (tokenCache.tokens === undefined) {
     try {
-      const tokens = window.localStorage.getItem("hassTokens");
+      const tokens = storage.getItem("hassTokens");
       if (tokens) {
         tokenCache.tokens = JSON.parse(tokens);
         tokenCache.writeEnabled = true;

src/util/ha-pref-storage.ts

@@ -1,5 +1,8 @@
+import { FallbackStorage } from "../../test/test_helper/local-storage-fallback";
 import type { HomeAssistant } from "../types";
 
+const storage = window.localStorage || new FallbackStorage();
+
 const STORED_STATE = [
   "dockedSidebar",
   "selectedTheme",
@@ -15,10 +18,7 @@ export function storeState(hass: HomeAssistant) {
   try {
     STORED_STATE.forEach((key) => {
       const value = hass[key];
-      window.localStorage.setItem(
-        key,
-        JSON.stringify(value === undefined ? null : value)
-      );
+      storage.setItem(key, JSON.stringify(value === undefined ? null : value));
     });
   } catch (err: any) {
     // Safari throws exception in private mode
@@ -35,7 +35,7 @@ export function getState() {
   const state = {};
 
   STORED_STATE.forEach((key) => {
-    const storageItem = window.localStorage.getItem(key);
+    const storageItem = storage.getItem(key);
     if (storageItem !== null) {
       let value = JSON.parse(storageItem);
       // selectedTheme went from string to object on 20200718
@@ -53,5 +53,5 @@ export function getState() {
 }
 
 export function clearState() {
-  window.localStorage.clear();
+  storage.clear();
 }

test/common/auth/token_storage/saveTokens.test.ts

@@ -58,9 +58,11 @@ describe("token_storage.saveTokens", () => {
       extractSearchParam: extractSearchParamSpy,
     }));
 
-    window.localStorage = {
-      setItem: vi.fn(),
-    } as unknown as Storage;
+    const { FallbackStorage: fallbackStorage } = await import(
+      "../../../test_helper/local-storage-fallback"
+    );
+    const setItemSpy = vi.fn();
+    fallbackStorage.prototype.setItem = setItemSpy;
 
     ({ saveTokens } = await import(
       "../../../../src/common/auth/token_storage"
@@ -70,8 +72,8 @@ describe("token_storage.saveTokens", () => {
     expect(window.__tokenCache.writeEnabled).toBe(true);
     expect(extractSearchParamSpy).toHaveBeenCalledOnce();
     expect(extractSearchParamSpy).toHaveBeenCalledWith("storeToken");
-    expect(window.localStorage.setItem).toHaveBeenCalledOnce();
-    expect(window.localStorage.setItem).toHaveBeenCalledWith(
+    expect(setItemSpy).toHaveBeenCalledOnce();
+    expect(setItemSpy).toHaveBeenCalledWith(
       "hassTokens",
       JSON.stringify(tokens)
     );
@@ -101,11 +103,13 @@ describe("token_storage.saveTokens", () => {
       extractSearchParam: extractSearchParamSpy,
     }));
 
-    window.localStorage = {
-      setItem: vi.fn(() => {
-        throw new Error("Full storage");
-      }),
-    } as unknown as Storage;
+    const { FallbackStorage: fallbackStorage } = await import(
+      "../../../test_helper/local-storage-fallback"
+    );
+    const setItemSpy = vi.fn(() => {
+      throw new Error("Full storage");
+    });
+    fallbackStorage.prototype.setItem = setItemSpy;
 
     // eslint-disable-next-line no-global-assign
     console = {
@@ -120,8 +124,8 @@ describe("token_storage.saveTokens", () => {
     expect(window.__tokenCache.tokens).toEqual(tokens);
     expect(window.__tokenCache.writeEnabled).toBe(true);
     expect(extractSearchParamSpy).toBeCalledTimes(0);
-    expect(window.localStorage.setItem).toHaveBeenCalledOnce();
-    expect(window.localStorage.setItem).toHaveBeenCalledWith(
+    expect(setItemSpy).toHaveBeenCalledOnce();
+    expect(setItemSpy).toHaveBeenCalledWith(
       "hassTokens",
       JSON.stringify(tokens)
     );

test/common/auth/token_storage/token_storage.test.ts

@@ -13,6 +13,7 @@ describe("token_storage", () => {
   });
 
   afterEach(() => {
+    vi.resetAllMocks();
     vi.resetModules();
   });
 
@@ -40,9 +41,11 @@ describe("token_storage", () => {
       clientId: "client",
     };
 
-    window.localStorage = {
-      getItem: vi.fn(() => JSON.stringify(tokens)),
-    } as unknown as Storage;
+    const { FallbackStorage: fallbackStorage } = await import(
+      "../../../test_helper/local-storage-fallback"
+    );
+    const getItemSpy = vi.fn(() => JSON.stringify(tokens));
+    fallbackStorage.prototype.getItem = getItemSpy;
 
     const { loadTokens } = await import(
       "../../../../src/common/auth/token_storage"
@@ -53,14 +56,16 @@ describe("token_storage", () => {
 
     expect(window.__tokenCache.tokens).toEqual(tokens);
     expect(window.__tokenCache.writeEnabled).toBe(true);
-    expect(window.localStorage.getItem).toHaveBeenCalledOnce();
-    expect(window.localStorage.getItem).toHaveBeenCalledWith("hassTokens");
+    expect(getItemSpy).toHaveBeenCalledOnce();
+    expect(getItemSpy).toHaveBeenCalledWith("hassTokens");
   });
 
   test("should load null tokens", async () => {
-    window.localStorage = {
-      getItem: vi.fn(() => null),
-    } as unknown as Storage;
+    const { FallbackStorage: fallbackStorage } = await import(
+      "../../../test_helper/local-storage-fallback"
+    );
+    const getItemSpy = vi.fn(() => "hello");
+    fallbackStorage.prototype.getItem = getItemSpy;
 
     const { loadTokens } = await import(
       "../../../../src/common/auth/token_storage"
@@ -71,8 +76,8 @@ describe("token_storage", () => {
 
     expect(window.__tokenCache.tokens).toEqual(null);
     expect(window.__tokenCache.writeEnabled).toBe(undefined);
-    expect(window.localStorage.getItem).toHaveBeenCalledOnce();
-    expect(window.localStorage.getItem).toHaveBeenCalledWith("hassTokens");
+    expect(getItemSpy).toHaveBeenCalledOnce();
+    expect(getItemSpy).toHaveBeenCalledWith("hassTokens");
   });
 
   it("should enable write", async () => {
@@ -91,18 +96,20 @@ describe("token_storage", () => {
       })
     );
 
-    window.localStorage = {
-      setItem: vi.fn(),
-    } as unknown as Storage;
+    const { FallbackStorage: fallbackStorage } = await import(
+      "../../../test_helper/local-storage-fallback"
+    );
+    const setItemSpy = vi.fn();
+    fallbackStorage.prototype.setItem = setItemSpy;
 
     const { enableWrite } = await import(
       "../../../../src/common/auth/token_storage"
     );
 
     enableWrite();
     expect(window.__tokenCache.writeEnabled).toBe(true);
-    expect(window.localStorage.setItem).toHaveBeenCalledOnce();
-    expect(window.localStorage.setItem).toHaveBeenCalledWith(
+    expect(setItemSpy).toHaveBeenCalledOnce();
+    expect(setItemSpy).toHaveBeenCalledWith(
       "hassTokens",
       JSON.stringify("testToken")
     );

test/test_helper/local-storage-fallback.ts

@@ -0,0 +1,38 @@
+export class FallbackStorage implements Storage {
+  private valuesMap = new Map();
+
+  getItem(key) {
+    const stringKey = String(key);
+    if (this.valuesMap.has(key)) {
+      return String(this.valuesMap.get(stringKey));
+    }
+    return null;
+  }
+
+  setItem(key, val) {
+    this.valuesMap.set(String(key), String(val));
+  }
+
+  removeItem(key) {
+    this.valuesMap.delete(key);
+  }
+
+  clear() {
+    this.valuesMap.clear();
+  }
+
+  key(i) {
+    if (arguments.length === 0) {
+      // this is a TypeError implemented on Chrome, Firefox throws Not enough arguments to Storage.key.
+      throw new TypeError(
+        "Failed to execute 'key' on 'Storage': 1 argument required, but only 0 present."
+      );
+    }
+    const arr = Array.from(this.valuesMap.keys());
+    return arr[i];
+  }
+
+  get length() {
+    return this.valuesMap.size;
+  }
+}