ai-details.test.ts206 lines · main
1import { beforeEach, describe, expect, it, vi } from 'vitest'
2
3import { getOrgAIDetails, getProjectAIDetails } from './ai-details'
4
5vi.mock('@/data/organizations/organizations-query', () => ({
6 getOrganizations: vi.fn(),
7}))
8
9vi.mock('@/data/projects/project-detail-query', () => ({
10 getProjectDetail: vi.fn(),
11}))
12
13vi.mock('@/data/subscriptions/org-subscription-query', () => ({
14 getOrgSubscription: vi.fn(),
15}))
16
17vi.mock('@/data/config/project-settings-v2-query', () => ({
18 getProjectSettings: vi.fn(),
19}))
20
21vi.mock('@/hooks/misc/useOrgOptedIntoAi', () => ({
22 getAiOptInLevel: vi.fn(),
23}))
24
25vi.mock('@/components/interfaces/Billing/Subscription/Subscription.utils', () => ({
26 subscriptionHasHipaaAddon: vi.fn(),
27}))
28
29vi.mock('@/data/entitlements/entitlements-query', () => ({
30 checkEntitlement: vi.fn(),
31}))
32
33const AUTH = 'Bearer token'
34const HEADERS = { 'Content-Type': 'application/json', Authorization: AUTH }
35
36describe('getOrgAIDetails', () => {
37 let mockGetOrganizations: ReturnType<typeof vi.fn>
38 let mockGetOrgSubscription: ReturnType<typeof vi.fn>
39 let mockGetAiOptInLevel: ReturnType<typeof vi.fn>
40 let mockSubscriptionHasHipaaAddon: ReturnType<typeof vi.fn>
41 let mockCheckEntitlement: ReturnType<typeof vi.fn>
42
43 beforeEach(async () => {
44 vi.clearAllMocks()
45
46 const orgsQuery = await import('@/data/organizations/organizations-query')
47 const subscriptionQuery = await import('@/data/subscriptions/org-subscription-query')
48 const aiHook = await import('@/hooks/misc/useOrgOptedIntoAi')
49 const subscriptionUtils =
50 await import('@/components/interfaces/Billing/Subscription/Subscription.utils')
51 const entitlementsQuery = await import('@/data/entitlements/entitlements-query')
52
53 mockGetOrganizations = vi.mocked(orgsQuery.getOrganizations)
54 mockGetOrgSubscription = vi.mocked(subscriptionQuery.getOrgSubscription)
55 mockGetAiOptInLevel = vi.mocked(aiHook.getAiOptInLevel)
56 mockSubscriptionHasHipaaAddon = vi.mocked(subscriptionUtils.subscriptionHasHipaaAddon)
57 mockCheckEntitlement = vi.mocked(entitlementsQuery.checkEntitlement)
58
59 mockGetOrgSubscription.mockResolvedValue({ addons: [] })
60 mockSubscriptionHasHipaaAddon.mockReturnValue(false)
61 mockCheckEntitlement.mockResolvedValue({ hasAccess: false })
62 })
63
64 it('returns org-level fields', async () => {
65 mockGetOrganizations.mockResolvedValue([
66 { id: 1, slug: 'test-org', plan: { id: 'pro' }, opt_in_tags: [] },
67 ])
68 mockGetAiOptInLevel.mockReturnValue('schema')
69
70 const result = await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
71
72 expect(result).toEqual({
73 aiOptInLevel: 'schema',
74 hasAccessToAdvanceModel: false,
75 hasHipaaAddon: false,
76 orgId: 1,
77 planId: 'pro',
78 })
79 })
80
81 it('returns hasAccessToAdvanceModel true when entitlement grants access', async () => {
82 mockGetOrganizations.mockResolvedValue([
83 { id: 1, slug: 'test-org', plan: { id: 'pro' }, opt_in_tags: [] },
84 ])
85 mockGetAiOptInLevel.mockReturnValue('schema')
86 mockCheckEntitlement.mockResolvedValue({ hasAccess: true })
87
88 const result = await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
89
90 expect(result.hasAccessToAdvanceModel).toBe(true)
91 })
92
93 it('returns hasHipaaAddon from subscription', async () => {
94 mockGetOrganizations.mockResolvedValue([
95 { id: 1, slug: 'test-org', plan: { id: 'enterprise' }, opt_in_tags: [] },
96 ])
97 mockGetAiOptInLevel.mockReturnValue('schema')
98 mockSubscriptionHasHipaaAddon.mockReturnValue(true)
99
100 const result = await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
101
102 expect(result.hasHipaaAddon).toBe(true)
103 })
104
105 it('calls getAiOptInLevel with the matched org opt_in_tags', async () => {
106 const opt_in_tags = ['AI_SQL_GENERATOR_OPT_IN']
107 mockGetOrganizations.mockResolvedValue([
108 { id: 1, slug: 'test-org', plan: { id: 'pro' }, opt_in_tags },
109 ])
110 mockGetAiOptInLevel.mockReturnValue('schema')
111
112 await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
113
114 expect(mockGetAiOptInLevel).toHaveBeenCalledWith(opt_in_tags)
115 })
116
117 it('forwards authorization headers to all fetches', async () => {
118 mockGetOrganizations.mockResolvedValue([
119 { id: 1, slug: 'test-org', plan: { id: 'pro' }, opt_in_tags: [] },
120 ])
121 mockGetAiOptInLevel.mockReturnValue('schema')
122
123 await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
124
125 expect(mockGetOrganizations).toHaveBeenCalledWith({ headers: HEADERS })
126 expect(mockGetOrgSubscription).toHaveBeenCalledWith({ orgSlug: 'test-org' }, undefined, HEADERS)
127 })
128
129 it('finds the correct org when multiple orgs are returned', async () => {
130 mockGetOrganizations.mockResolvedValue([
131 { id: 1, slug: 'org-1', plan: { id: 'free' }, opt_in_tags: [] },
132 { id: 2, slug: 'test-org', plan: { id: 'pro' }, opt_in_tags: [] },
133 ])
134 mockGetAiOptInLevel.mockReturnValue('schema')
135
136 const result = await getOrgAIDetails({ orgSlug: 'test-org', authorization: AUTH })
137
138 expect(result.orgId).toBe(2)
139 expect(result.planId).toBe('pro')
140 })
141})
142
143describe('getProjectAIDetails', () => {
144 let mockGetProjectDetail: ReturnType<typeof vi.fn>
145 let mockGetProjectSettings: ReturnType<typeof vi.fn>
146
147 beforeEach(async () => {
148 vi.clearAllMocks()
149
150 const projectQuery = await import('@/data/projects/project-detail-query')
151 const settingsQuery = await import('@/data/config/project-settings-v2-query')
152
153 mockGetProjectDetail = vi.mocked(projectQuery.getProjectDetail)
154 mockGetProjectSettings = vi.mocked(settingsQuery.getProjectSettings)
155 })
156
157 it('returns region and isSensitive', async () => {
158 mockGetProjectDetail.mockResolvedValue({ ref: 'test-project', region: 'us-east-1' })
159 mockGetProjectSettings.mockResolvedValue({ is_sensitive: false })
160
161 const result = await getProjectAIDetails({ projectRef: 'test-project', authorization: AUTH })
162
163 expect(result).toEqual({ region: 'us-east-1', isSensitive: false })
164 })
165
166 it('returns isSensitive true when project is marked sensitive', async () => {
167 mockGetProjectDetail.mockResolvedValue({ ref: 'test-project', region: 'us-east-1' })
168 mockGetProjectSettings.mockResolvedValue({ is_sensitive: true })
169
170 const result = await getProjectAIDetails({ projectRef: 'test-project', authorization: AUTH })
171
172 expect(result.isSensitive).toBe(true)
173 })
174
175 it('returns isSensitive undefined when project settings are unavailable', async () => {
176 mockGetProjectDetail.mockResolvedValue({ ref: 'test-project', region: 'us-east-1' })
177 mockGetProjectSettings.mockResolvedValue(undefined)
178
179 const result = await getProjectAIDetails({ projectRef: 'test-project', authorization: AUTH })
180
181 expect(result.isSensitive).toBeUndefined()
182 })
183
184 it('returns region undefined when project detail is unavailable', async () => {
185 mockGetProjectDetail.mockResolvedValue(undefined)
186 mockGetProjectSettings.mockResolvedValue({ is_sensitive: false })
187
188 const result = await getProjectAIDetails({ projectRef: 'test-project', authorization: AUTH })
189
190 expect(result.region).toBeUndefined()
191 })
192
193 it('forwards authorization headers to all fetches', async () => {
194 mockGetProjectDetail.mockResolvedValue({ ref: 'test-project', region: 'us-east-1' })
195 mockGetProjectSettings.mockResolvedValue({ is_sensitive: false })
196
197 await getProjectAIDetails({ projectRef: 'test-project', authorization: AUTH })
198
199 expect(mockGetProjectDetail).toHaveBeenCalledWith({ ref: 'test-project' }, undefined, HEADERS)
200 expect(mockGetProjectSettings).toHaveBeenCalledWith(
201 { projectRef: 'test-project' },
202 undefined,
203 HEADERS
204 )
205 })
206})