Refac: Centralize storage file management (#4078)

Co-authored-by: Taylor Mullen <[email protected]>

1 files changed, 339 insertions, 0 deletions

diff --git a/packages/core/src/utils/userAccountManager.test.ts b/packages/core/src/utils/userAccountManager.test.ts
new file mode 100644
index 00000000..1e7e05aa
--- /dev/null
+++ b/packages/core/src/utils/userAccountManager.test.ts
@@ -0,0 +1,339 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { vi, describe, it, expect, beforeEach, afterEach, Mock } from 'vitest';
+import { UserAccountManager } from './userAccountManager.js';
+import * as fs from 'node:fs';
+import * as os from 'node:os';
+import path from 'node:path';
+
+vi.mock('os', async (importOriginal) => {
+  const os = await importOriginal<typeof import('os')>();
+  return {
+    ...os,
+    homedir: vi.fn(),
+  };
+});
+
+describe('UserAccountManager', () => {
+  let tempHomeDir: string;
+  let userAccountManager: UserAccountManager;
+  let accountsFile: () => string;
+
+  beforeEach(() => {
+    tempHomeDir = fs.mkdtempSync(
+      path.join(os.tmpdir(), 'gemini-cli-test-home-'),
+    );
+    (os.homedir as Mock).mockReturnValue(tempHomeDir);
+    accountsFile = () =>
+      path.join(tempHomeDir, '.gemini', 'google_accounts.json');
+    userAccountManager = new UserAccountManager();
+  });
+
+  afterEach(() => {
+    fs.rmSync(tempHomeDir, { recursive: true, force: true });
+    vi.clearAllMocks();
+  });
+
+  describe('cacheGoogleAccount', () => {
+    it('should create directory and write initial account file', async () => {
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+
+      // Verify Google Account ID was cached
+      expect(fs.existsSync(accountsFile())).toBe(true);
+      expect(fs.readFileSync(accountsFile(), 'utf-8')).toBe(
+        JSON.stringify({ active: '[email protected]', old: [] }, null, 2),
+      );
+    });
+
+    it('should update active account and move previous to old', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify(
+          { active: '[email protected]', old: ['[email protected]'] },
+          null,
+          2,
+        ),
+      );
+
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+
+      expect(fs.readFileSync(accountsFile(), 'utf-8')).toBe(
+        JSON.stringify(
+          {
+            active: '[email protected]',
+            old: ['[email protected]', '[email protected]'],
+          },
+          null,
+          2,
+        ),
+      );
+    });
+
+    it('should not add a duplicate to the old list', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify(
+          { active: '[email protected]', old: ['[email protected]'] },
+          null,
+          2,
+        ),
+      );
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+
+      expect(fs.readFileSync(accountsFile(), 'utf-8')).toBe(
+        JSON.stringify(
+          { active: '[email protected]', old: ['[email protected]'] },
+          null,
+          2,
+        ),
+      );
+    });
+
+    it('should handle corrupted JSON by starting fresh', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), 'not valid json');
+      const consoleLogSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+
+      expect(consoleLogSpy).toHaveBeenCalled();
+      expect(JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'))).toEqual({
+        active: '[email protected]',
+        old: [],
+      });
+    });
+
+    it('should handle valid JSON with incorrect schema by starting fresh', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({ active: '[email protected]', old: 'not-an-array' }),
+      );
+      const consoleLogSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      await userAccountManager.cacheGoogleAccount('[email protected]');
+
+      expect(consoleLogSpy).toHaveBeenCalled();
+      expect(JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'))).toEqual({
+        active: '[email protected]',
+        old: [],
+      });
+    });
+  });
+
+  describe('getCachedGoogleAccount', () => {
+    it('should return the active account if file exists and is valid', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({ active: '[email protected]', old: [] }, null, 2),
+      );
+      const account = userAccountManager.getCachedGoogleAccount();
+      expect(account).toBe('[email protected]');
+    });
+
+    it('should return null if file does not exist', () => {
+      const account = userAccountManager.getCachedGoogleAccount();
+      expect(account).toBeNull();
+    });
+
+    it('should return null if file is empty', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), '');
+      const account = userAccountManager.getCachedGoogleAccount();
+      expect(account).toBeNull();
+    });
+
+    it('should return null and log if file is corrupted', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), '{ "active": "[email protected]"'); // Invalid JSON
+      const consoleLogSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      const account = userAccountManager.getCachedGoogleAccount();
+
+      expect(account).toBeNull();
+      expect(consoleLogSpy).toHaveBeenCalled();
+    });
+
+    it('should return null if active key is missing', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), JSON.stringify({ old: [] }));
+      const account = userAccountManager.getCachedGoogleAccount();
+      expect(account).toBeNull();
+    });
+  });
+
+  describe('clearCachedGoogleAccount', () => {
+    it('should set active to null and move it to old', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify(
+          { active: '[email protected]', old: ['[email protected]'] },
+          null,
+          2,
+        ),
+      );
+
+      await userAccountManager.clearCachedGoogleAccount();
+
+      const stored = JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'));
+      expect(stored.active).toBeNull();
+      expect(stored.old).toEqual(['[email protected]', '[email protected]']);
+    });
+
+    it('should handle empty file gracefully', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), '');
+      await userAccountManager.clearCachedGoogleAccount();
+      const stored = JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'));
+      expect(stored.active).toBeNull();
+      expect(stored.old).toEqual([]);
+    });
+
+    it('should handle corrupted JSON by creating a fresh file', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), 'not valid json');
+      const consoleLogSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      await userAccountManager.clearCachedGoogleAccount();
+
+      expect(consoleLogSpy).toHaveBeenCalled();
+      const stored = JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'));
+      expect(stored.active).toBeNull();
+      expect(stored.old).toEqual([]);
+    });
+
+    it('should be idempotent if active account is already null', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({ active: null, old: ['[email protected]'] }, null, 2),
+      );
+
+      await userAccountManager.clearCachedGoogleAccount();
+
+      const stored = JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'));
+      expect(stored.active).toBeNull();
+      expect(stored.old).toEqual(['[email protected]']);
+    });
+
+    it('should not add a duplicate to the old list', async () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify(
+          {
+            active: '[email protected]',
+            old: ['[email protected]'],
+          },
+          null,
+          2,
+        ),
+      );
+
+      await userAccountManager.clearCachedGoogleAccount();
+
+      const stored = JSON.parse(fs.readFileSync(accountsFile(), 'utf-8'));
+      expect(stored.active).toBeNull();
+      expect(stored.old).toEqual(['[email protected]']);
+    });
+  });
+
+  describe('getLifetimeGoogleAccounts', () => {
+    it('should return 0 if the file does not exist', () => {
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(0);
+    });
+
+    it('should return 0 if the file is empty', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), '');
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(0);
+    });
+
+    it('should return 0 if the file is corrupted', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(accountsFile(), 'invalid json');
+      const consoleDebugSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(0);
+      expect(consoleDebugSpy).toHaveBeenCalled();
+    });
+
+    it('should return 1 if there is only an active account', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({ active: '[email protected]', old: [] }),
+      );
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(1);
+    });
+
+    it('should correctly count old accounts when active is null', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({
+          active: null,
+          old: ['[email protected]', '[email protected]'],
+        }),
+      );
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(2);
+    });
+
+    it('should correctly count both active and old accounts', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({
+          active: '[email protected]',
+          old: ['[email protected]', '[email protected]'],
+        }),
+      );
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(3);
+    });
+
+    it('should handle valid JSON with incorrect schema by returning 0', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({ active: null, old: 1 }),
+      );
+      const consoleLogSpy = vi
+        .spyOn(console, 'log')
+        .mockImplementation(() => {});
+
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(0);
+      expect(consoleLogSpy).toHaveBeenCalled();
+    });
+
+    it('should not double count if active account is also in old list', () => {
+      fs.mkdirSync(path.dirname(accountsFile()), { recursive: true });
+      fs.writeFileSync(
+        accountsFile(),
+        JSON.stringify({
+          active: '[email protected]',
+          old: ['[email protected]', '[email protected]'],
+        }),
+      );
+      expect(userAccountManager.getLifetimeGoogleAccounts()).toBe(2);
+    });
+  });
+});