Fix coverage

Author: Mrtenz

packages/snaps-controllers/coverage.json

@@ -1,6 +1,6 @@
 {
-  "branches": 93.34,
-  "functions": 97.37,
-  "lines": 98.33,
+  "branches": 93.36,
+  "functions": 97.38,
+  "lines": 98.34,
   "statements": 98.07
 }

packages/snaps-controllers/src/snaps/SnapController.test.tsx

@@ -9736,10 +9736,27 @@ describe('SnapController', () => {
       expect(encryptWithKey).not.toHaveBeenCalled();
 
       jest.advanceTimersByTime(5_000);
-
       await promise;
+
       expect(encryptWithKey).toHaveBeenCalledTimes(1);
 
+      const nextStateChange = waitForStateChange(messenger);
+      await messenger.call(
+        'SnapController:updateSnapState',
+        MOCK_SNAP_ID,
+        { qux: 'quux' },
+        true,
+      );
+
+      expect(
+        await messenger.call('SnapController:getSnapState', MOCK_SNAP_ID, true),
+      ).toStrictEqual({ qux: 'quux' });
+
+      jest.advanceTimersByTime(5_000);
+      await nextStateChange;
+
+      expect(encryptWithKey).toHaveBeenCalledTimes(2);
+
       snapController.destroy();
     });
 

packages/snaps-controllers/src/snaps/SnapController.ts

@@ -1799,6 +1799,23 @@ export class SnapController extends BaseController<
     return truncateSnap(this.getExpect(snapId));
   }
 
+  /**
+   * Check if a given Snap has a cached encryption key stored in the runtime.
+   *
+   * @param snapId - The Snap ID.
+   * @param runtime - The Snap runtime data.
+   * @returns True if the Snap has a cached encryption key, otherwise false.
+   */
+  #hasCachedEncryptionKey(
+    snapId: SnapId,
+    runtime = this.#getRuntimeExpect(snapId),
+  ): runtime is SnapRuntimeData & {
+    encryptionKey: string;
+    encryptionSalt: string;
+  } {
+    return runtime.encryptionKey !== null && runtime.encryptionSalt !== null;
+  }
+
   /**
    * Generate an encryption key to be used for state encryption for a given Snap.
    *
@@ -1822,7 +1839,7 @@ export class SnapController extends BaseController<
   }): Promise<{ key: unknown; salt: string }> {
     const runtime = this.#getRuntimeExpect(snapId);
 
-    if (runtime.encryptionKey && runtime.encryptionSalt && useCache) {
+    if (this.#hasCachedEncryptionKey(snapId, runtime) && useCache) {
       return {
         key: await this.#encryptor.importKey(runtime.encryptionKey),
         salt: runtime.encryptionSalt,
@@ -1854,17 +1871,6 @@ export class SnapController extends BaseController<
     return { key: encryptionKey, salt };
   }
 
-  /**
-   * Check if a given Snap has a cached encryption key stored in the runtime.
-   *
-   * @param snapId - The Snap ID.
-   * @returns True if the Snap has a cached encryption key, otherwise false.
-   */
-  #hasCachedEncryptionKey(snapId: SnapId) {
-    const runtime = this.#getRuntimeExpect(snapId);
-    return runtime.encryptionKey !== null && runtime.encryptionSalt !== null;
-  }
-
   /**
    * Decrypt the encrypted state for a given Snap.
    *