feat: add backup & sync capabilities to AccountTreeController

packages/account-tree-controller/src/AccountTreeController.test.ts

@@ -23,7 +23,10 @@ import { KeyringTypes } from '@metamask/keyring-controller';
 import type { InternalAccount } from '@metamask/keyring-internal-api';
 import type { GetSnap as SnapControllerGetSnap } from '@metamask/snaps-controllers';
 
-import { AccountTreeController } from './AccountTreeController';
+import {
+  AccountTreeController,
+  getDefaultAccountTreeControllerState,
+} from './AccountTreeController';
 import type { BackupAndSyncAnalyticsEventPayload } from './backup-and-sync/analytics';
 import { BackupAndSyncService } from './backup-and-sync/service';
 import { getAccountWalletNameFromKeyringType } from './rules/keyring';
@@ -225,6 +228,7 @@ function getAccountTreeControllerMessenger(
       'AccountsController:getAccount',
       'AccountsController:getSelectedAccount',
       'AccountsController:setSelectedAccount',
+      'UserStorageController:getState',
       'UserStorageController:performGetStorage',
       'UserStorageController:performGetStorageAllFeatureEntries',
       'UserStorageController:performSetStorage',
@@ -248,15 +252,22 @@ function getAccountTreeControllerMessenger(
  * @param options.config - Configuration options for the controller.
  * @param options.config.backupAndSync - Configuration options for backup and sync.
  * @param options.config.backupAndSync.onBackupAndSyncEvent - Event handler for backup and sync events.
- * @param options.config.backupAndSync.enableDebugLogging - Flag to enable debug logging.
+ * @param options.config.backupAndSync.isAccountSyncingEnabled - Flag to enable account syncing.
+ * @param options.config.backupAndSync.isBackupAndSyncEnabled - Flag to enable backup and sync.
  * @returns An object containing the controller instance and the messenger.
  */
 function setup({
   state = {},
   messenger = getRootMessenger(),
   accounts = [],
   keyrings = [],
-  config,
+  config = {
+    backupAndSync: {
+      isAccountSyncingEnabled: true,
+      isBackupAndSyncEnabled: true,
+      onBackupAndSyncEvent: jest.fn(),
+    },
+  },
 }: {
   state?: Partial<AccountTreeControllerState>;
   messenger?: Messenger<
@@ -267,10 +278,11 @@ function setup({
   keyrings?: KeyringObject[];
   config?: {
     backupAndSync?: {
+      isAccountSyncingEnabled?: boolean;
+      isBackupAndSyncEnabled?: boolean;
       onBackupAndSyncEvent?: (
         event: BackupAndSyncAnalyticsEventPayload,
       ) => void;
-      enableDebugLogging?: boolean;
     };
   };
 } = {}): {
@@ -317,6 +329,7 @@ function setup({
       getSelectedAccount: jest.fn(),
     },
     UserStorageController: {
+      getState: jest.fn(),
       performGetStorage: jest.fn(),
       performGetStorageAllFeatureEntries: jest.fn(),
       performSetStorage: jest.fn(),
@@ -372,6 +385,15 @@ function setup({
     );
 
     // Mock UserStorageController methods
+    mocks.UserStorageController.getState.mockImplementation(() => ({
+      isBackupAndSyncEnabled: config?.backupAndSync?.isBackupAndSyncEnabled,
+      isAccountSyncingEnabled: config?.backupAndSync?.isAccountSyncingEnabled,
+    }));
+    messenger.registerActionHandler(
+      'UserStorageController:getState',
+      mocks.UserStorageController.getState,
+    );
+
     messenger.registerActionHandler(
       'UserStorageController:performGetStorage',
       mocks.UserStorageController.performGetStorage,
@@ -2833,11 +2855,12 @@ describe('AccountTreeController', () => {
       jest.clearAllMocks();
     });
 
-    it('calls performFullSync on the syncing service', async () => {
+    it('calls fullSync on the syncing service', async () => {
       // Spy on the BackupAndSyncService constructor and methods
-      const performFullSyncSpy = jest
-        .spyOn(BackupAndSyncService.prototype, 'performFullSync')
-        .mockResolvedValue(undefined);
+      const fullSyncSpy = jest.spyOn(
+        BackupAndSyncService.prototype,
+        'performFullSync',
+      );
 
       const { controller } = setup({
         accounts: [MOCK_HARDWARE_ACCOUNT_1], // Use hardware account to avoid entropy calls
@@ -2848,12 +2871,12 @@ describe('AccountTreeController', () => {
 
       await controller.syncWithUserStorage();
 
-      expect(performFullSyncSpy).toHaveBeenCalledTimes(1);
+      expect(fullSyncSpy).toHaveBeenCalledTimes(1);
     });
 
     it('handles sync errors gracefully', async () => {
       const syncError = new Error('Sync failed');
-      const performFullSyncSpy = jest
+      const fullSyncSpy = jest
         .spyOn(BackupAndSyncService.prototype, 'performFullSync')
         .mockRejectedValue(syncError);
 
@@ -2867,7 +2890,7 @@ describe('AccountTreeController', () => {
       await expect(controller.syncWithUserStorage()).rejects.toThrow(
         syncError.message,
       );
-      expect(performFullSyncSpy).toHaveBeenCalledTimes(1);
+      expect(fullSyncSpy).toHaveBeenCalledTimes(1);
     });
   });
 
@@ -2895,28 +2918,27 @@ describe('AccountTreeController', () => {
       expect(controller.state.accountWalletsMetadata).not.toStrictEqual({});
 
       // Clear the metadata
-      controller.clearPersistedMetadataAndSyncingState();
+      controller.clearState();
 
       // Verify everything is cleared
-      expect(controller.state.accountGroupsMetadata).toStrictEqual({});
-      expect(controller.state.accountWalletsMetadata).toStrictEqual({});
-      expect(controller.state.hasAccountTreeSyncingSyncedAtLeastOnce).toBe(
-        false,
+      expect(controller.state).toStrictEqual(
+        getDefaultAccountTreeControllerState(),
       );
     });
   });
 
   describe('backup and sync config initialization', () => {
-    it('initializes backup and sync config with provided analytics callback and debug logging', async () => {
+    it('initializes backup and sync config with provided analytics callback', async () => {
       const mockAnalyticsCallback = jest.fn();
 
       const { controller } = setup({
         accounts: [MOCK_HD_ACCOUNT_1],
         keyrings: [MOCK_HD_KEYRING_1],
         config: {
           backupAndSync: {
+            isAccountSyncingEnabled: true,
+            isBackupAndSyncEnabled: true,
             onBackupAndSyncEvent: mockAnalyticsCallback,
-            enableDebugLogging: true,
           },
         },
       });

packages/account-tree-controller/src/AccountTreeController.ts

@@ -1110,19 +1110,15 @@ export class AccountTreeController extends BaseController<
   }
 
   /**
-   * Clears all persisted metadata and syncing state.
-   *
-   * This will reset the account groups and wallets metadata, as well as
-   * the syncing state. It also calls `init()` to reinitialize the controller.
-   * This should be used when we want to completely reset the controller's state.
+   * Clears the controller state and resets to default values.
    */
-  clearPersistedMetadataAndSyncingState(): void {
-    this.update((state) => {
-      state.accountGroupsMetadata = {};
-      state.accountWalletsMetadata = {};
-      state.hasAccountTreeSyncingSyncedAtLeastOnce = false;
+  clearState(): void {
+    this.update(() => {
+      return {
+        ...getDefaultAccountTreeControllerState(),
+      };
     });
-    this.init();
+    this.#backupAndSyncService.clearState();
   }
 
   /**

packages/account-tree-controller/src/backup-and-sync/service/atomic-sync-queue.test.ts

@@ -1,6 +1,5 @@
 import { AtomicSyncQueue } from './atomic-sync-queue';
 import { backupAndSyncLogger } from '../../logger';
-import type { BackupAndSyncContext } from '../types';
 
 jest.mock('../../logger', () => ({
   backupAndSyncLogger: jest.fn(),
@@ -12,18 +11,10 @@ const mockBackupAndSyncLogger = backupAndSyncLogger as jest.MockedFunction<
 
 describe('BackupAndSync - Service - AtomicSyncQueue', () => {
   let atomicSyncQueue: AtomicSyncQueue;
-  const mockContext = {
-    controller: {
-      state: {
-        isAccountTreeSyncingInProgress: false,
-        hasAccountTreeSyncingSyncedAtLeastOnce: true,
-      },
-    },
-  } as unknown as BackupAndSyncContext;
 
   beforeEach(() => {
     jest.clearAllMocks();
-    atomicSyncQueue = new AtomicSyncQueue(mockContext);
+    atomicSyncQueue = new AtomicSyncQueue();
   });
 
   afterEach(() => {
@@ -32,13 +23,13 @@ describe('BackupAndSync - Service - AtomicSyncQueue', () => {
 
   describe('constructor', () => {
     it('initializes with default debug logging function', () => {
-      const queue = new AtomicSyncQueue(mockContext);
+      const queue = new AtomicSyncQueue();
       expect(queue.size).toBe(0);
       expect(queue.isProcessing).toBe(false);
     });
 
     it('initializes with provided debug logging function', () => {
-      const queue = new AtomicSyncQueue(mockContext);
+      const queue = new AtomicSyncQueue();
       expect(queue.size).toBe(0);
       expect(queue.isProcessing).toBe(false);
     });
@@ -53,52 +44,6 @@ describe('BackupAndSync - Service - AtomicSyncQueue', () => {
       expect(atomicSyncQueue.size).toBe(1);
     });
 
-    it('does not enqueue when big sync is in progress', () => {
-      const mockContextWithBigSyncInProgress = {
-        ...mockContext,
-        controller: {
-          ...mockContext.controller,
-          state: {
-            ...mockContext.controller.state,
-            isAccountTreeSyncingInProgress: true,
-          },
-        },
-      } as unknown as BackupAndSyncContext;
-
-      const atomicSyncQueueWithBigSyncInProgress = new AtomicSyncQueue(
-        mockContextWithBigSyncInProgress,
-      );
-
-      const mockSyncFunction = jest.fn().mockResolvedValue(undefined);
-
-      atomicSyncQueueWithBigSyncInProgress.enqueue(mockSyncFunction);
-
-      expect(atomicSyncQueueWithBigSyncInProgress.size).toBe(0);
-    });
-
-    it('does not enqueue if big sync has never been ran', () => {
-      const mockContextWithNoBigSyncEver = {
-        ...mockContext,
-        controller: {
-          ...mockContext.controller,
-          state: {
-            ...mockContext.controller.state,
-            hasAccountTreeSyncingSyncedAtLeastOnce: false,
-          },
-        },
-      } as unknown as BackupAndSyncContext;
-
-      const atomicSyncQueueWithNoBigSyncEver = new AtomicSyncQueue(
-        mockContextWithNoBigSyncEver,
-      );
-
-      const mockSyncFunction = jest.fn().mockResolvedValue(undefined);
-
-      atomicSyncQueueWithNoBigSyncEver.enqueue(mockSyncFunction);
-
-      expect(atomicSyncQueueWithNoBigSyncEver.size).toBe(0);
-    });
-
     it('triggers async processing after enqueueing', async () => {
       jest.useFakeTimers();
       const mockSyncFunction = jest.fn().mockResolvedValue(undefined);
@@ -131,32 +76,6 @@ describe('BackupAndSync - Service - AtomicSyncQueue', () => {
       expect(atomicSyncQueue.size).toBe(0);
     });
 
-    it('does not process when big sync is in progress', async () => {
-      const mockContextWithBigSyncInProgress = {
-        ...mockContext,
-        controller: {
-          ...mockContext.controller,
-          state: {
-            ...mockContext.controller.state,
-            isAccountTreeSyncingInProgress: true,
-          },
-        },
-      } as unknown as BackupAndSyncContext;
-
-      const atomicSyncQueueWithBigSyncInProgress = new AtomicSyncQueue(
-        mockContextWithBigSyncInProgress,
-      );
-
-      const mockSyncFunction = jest.fn().mockResolvedValue(undefined);
-
-      atomicSyncQueueWithBigSyncInProgress.enqueue(mockSyncFunction);
-
-      await atomicSyncQueueWithBigSyncInProgress.process();
-
-      expect(mockSyncFunction).not.toHaveBeenCalled();
-      expect(atomicSyncQueueWithBigSyncInProgress.size).toBe(0);
-    });
-
     it('does not process when already processing', async () => {
       const mockSyncFunction = jest.fn().mockImplementation(async () => {
         // While first function is processing, try to process again
@@ -241,7 +160,7 @@ describe('BackupAndSync - Service - AtomicSyncQueue', () => {
 
     it('accesses size property correctly', () => {
       // Create a fresh queue to test size property
-      const freshQueue = new AtomicSyncQueue(mockContext);
+      const freshQueue = new AtomicSyncQueue();
       expect(freshQueue.size).toBe(0);
 
       // Add multiple items
@@ -275,6 +194,19 @@ describe('BackupAndSync - Service - AtomicSyncQueue', () => {
         error,
       );
     });
+
+    it('rejects promise when awaited sync function throws error', async () => {
+      const error = new Error('Sync function failed');
+      const mockSyncFunction = jest.fn().mockRejectedValue(error);
+
+      const promise = atomicSyncQueue.enqueue(mockSyncFunction, {
+        await: true,
+      });
+
+      await expect(promise).rejects.toThrow('Sync function failed');
+      expect(mockSyncFunction).toHaveBeenCalled();
+    });
+
     it('handles empty queue after shift operation', async () => {
       // Test the scenario where shift() might return undefined/null
       // This can happen in race conditions or edge cases