summaryrefslogtreecommitdiff
path: root/packages/core/src/code_assist/setup.test.ts
blob: c1260e3f681d17f8d4535735501de2a70f02f8d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { setupUser, ProjectIdRequiredError } from './setup.js';
import { CodeAssistServer } from '../code_assist/server.js';
import { OAuth2Client } from 'google-auth-library';
import { GeminiUserTier, UserTierId } from './types.js';

vi.mock('../code_assist/server.js');

const mockPaidTier: GeminiUserTier = {
  id: UserTierId.STANDARD,
  name: 'paid',
  description: 'Paid tier',
};

describe('setupUser', () => {
  let mockLoad: ReturnType<typeof vi.fn>;
  let mockOnboardUser: ReturnType<typeof vi.fn>;

  beforeEach(() => {
    vi.resetAllMocks();
    mockLoad = vi.fn();
    mockOnboardUser = vi.fn().mockResolvedValue({
      done: true,
      response: {
        cloudaicompanionProject: {
          id: 'server-project',
        },
      },
    });
    vi.mocked(CodeAssistServer).mockImplementation(
      () =>
        ({
          loadCodeAssist: mockLoad,
          onboardUser: mockOnboardUser,
        }) as unknown as CodeAssistServer,
    );
  });

  it('should use GOOGLE_CLOUD_PROJECT when set', async () => {
    process.env.GOOGLE_CLOUD_PROJECT = 'test-project';
    mockLoad.mockResolvedValue({
      currentTier: mockPaidTier,
    });
    await setupUser({} as OAuth2Client);
    expect(CodeAssistServer).toHaveBeenCalledWith(
      {},
      'test-project',
      {},
      '',
      undefined,
    );
  });

  it('should treat empty GOOGLE_CLOUD_PROJECT as undefined and use project from server', async () => {
    process.env.GOOGLE_CLOUD_PROJECT = '';
    mockLoad.mockResolvedValue({
      cloudaicompanionProject: 'server-project',
      currentTier: mockPaidTier,
    });
    const projectId = await setupUser({} as OAuth2Client);
    expect(CodeAssistServer).toHaveBeenCalledWith(
      {},
      undefined,
      {},
      '',
      undefined,
    );
    expect(projectId).toEqual({
      projectId: 'server-project',
      userTier: 'standard-tier',
    });
  });

  it('should throw ProjectIdRequiredError when no project ID is available', async () => {
    delete process.env.GOOGLE_CLOUD_PROJECT;
    // And the server itself requires a project ID internally
    vi.mocked(CodeAssistServer).mockImplementation(() => {
      throw new ProjectIdRequiredError();
    });

    await expect(setupUser({} as OAuth2Client)).rejects.toThrow(
      ProjectIdRequiredError,
    );
  });
});