Skip to content

Commit 4e6cb8f

Browse files
waleedlatif1waleedlatif1
committed
fix(knowledge): fix search embedding test mocks, parallelize billing lookups
1 parent 57f190a commit 4e6cb8f

File tree

2 files changed

+51
-64
lines changed

2 files changed

+51
-64
lines changed

‎apps/sim/app/api/knowledge/search/utils.test.ts‎

Lines changed: 43 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* @vitest-environment node
66
*/
77
import { createEnvMock, databaseMock, loggerMock } from '@sim/testing'
8+
import { mockNextFetchResponse } from '@sim/testing/mocks'
89
import { beforeEach, describe, expect, it, vi } from 'vitest'
910

1011
vi.mock('drizzle-orm')
@@ -14,16 +15,6 @@ vi.mock('@/lib/knowledge/documents/utils', () => ({
1415
retryWithExponentialBackoff: (fn: any) => fn(),
1516
}))
1617

17-
vi.stubGlobal(
18-
'fetch',
19-
vi.fn().mockResolvedValue({
20-
ok: true,
21-
json: async () => ({
22-
data: [{ embedding: [0.1, 0.2, 0.3] }],
23-
}),
24-
})
25-
)
26-
2718
vi.mock('@/lib/core/config/env', () => createEnvMock())
2819

2920
import {
@@ -178,17 +169,16 @@ describe('Knowledge Search Utils', () => {
178169
OPENAI_API_KEY: 'test-openai-key',
179170
})
180171

181-
const fetchSpy = vi.mocked(fetch)
182-
fetchSpy.mockResolvedValueOnce({
183-
ok: true,
184-
json: async () => ({
172+
mockNextFetchResponse({
173+
json: {
185174
data: [{ embedding: [0.1, 0.2, 0.3] }],
186-
}),
187-
} as any)
175+
usage: { prompt_tokens: 1, total_tokens: 1 },
176+
},
177+
})
188178

189179
const result = await generateSearchEmbedding('test query')
190180

191-
expect(fetchSpy).toHaveBeenCalledWith(
181+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
192182
'https://test.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-12-01-preview',
193183
expect.objectContaining({
194184
headers: expect.objectContaining({
@@ -209,17 +199,16 @@ describe('Knowledge Search Utils', () => {
209199
OPENAI_API_KEY: 'test-openai-key',
210200
})
211201

212-
const fetchSpy = vi.mocked(fetch)
213-
fetchSpy.mockResolvedValueOnce({
214-
ok: true,
215-
json: async () => ({
202+
mockNextFetchResponse({
203+
json: {
216204
data: [{ embedding: [0.1, 0.2, 0.3] }],
217-
}),
218-
} as any)
205+
usage: { prompt_tokens: 1, total_tokens: 1 },
206+
},
207+
})
219208

220209
const result = await generateSearchEmbedding('test query')
221210

222-
expect(fetchSpy).toHaveBeenCalledWith(
211+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
223212
'https://api.openai.com/v1/embeddings',
224213
expect.objectContaining({
225214
headers: expect.objectContaining({
@@ -243,17 +232,16 @@ describe('Knowledge Search Utils', () => {
243232
OPENAI_API_KEY: 'test-openai-key',
244233
})
245234

246-
const fetchSpy = vi.mocked(fetch)
247-
fetchSpy.mockResolvedValueOnce({
248-
ok: true,
249-
json: async () => ({
235+
mockNextFetchResponse({
236+
json: {
250237
data: [{ embedding: [0.1, 0.2, 0.3] }],
251-
}),
252-
} as any)
238+
usage: { prompt_tokens: 1, total_tokens: 1 },
239+
},
240+
})
253241

254242
await generateSearchEmbedding('test query')
255243

256-
expect(fetchSpy).toHaveBeenCalledWith(
244+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
257245
expect.stringContaining('api-version='),
258246
expect.any(Object)
259247
)
@@ -273,17 +261,16 @@ describe('Knowledge Search Utils', () => {
273261
OPENAI_API_KEY: 'test-openai-key',
274262
})
275263

276-
const fetchSpy = vi.mocked(fetch)
277-
fetchSpy.mockResolvedValueOnce({
278-
ok: true,
279-
json: async () => ({
264+
mockNextFetchResponse({
265+
json: {
280266
data: [{ embedding: [0.1, 0.2, 0.3] }],
281-
}),
282-
} as any)
267+
usage: { prompt_tokens: 1, total_tokens: 1 },
268+
},
269+
})
283270

284271
await generateSearchEmbedding('test query', 'text-embedding-3-small')
285272

286-
expect(fetchSpy).toHaveBeenCalledWith(
273+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
287274
'https://test.openai.azure.com/openai/deployments/custom-embedding-model/embeddings?api-version=2024-12-01-preview',
288275
expect.any(Object)
289276
)
@@ -311,13 +298,12 @@ describe('Knowledge Search Utils', () => {
311298
KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
312299
})
313300

314-
const fetchSpy = vi.mocked(fetch)
315-
fetchSpy.mockResolvedValueOnce({
301+
mockNextFetchResponse({
316302
ok: false,
317303
status: 404,
318304
statusText: 'Not Found',
319-
text: async () => 'Deployment not found',
320-
} as any)
305+
text: 'Deployment not found',
306+
})
321307

322308
await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
323309

@@ -332,13 +318,12 @@ describe('Knowledge Search Utils', () => {
332318
OPENAI_API_KEY: 'test-openai-key',
333319
})
334320

335-
const fetchSpy = vi.mocked(fetch)
336-
fetchSpy.mockResolvedValueOnce({
321+
mockNextFetchResponse({
337322
ok: false,
338323
status: 429,
339324
statusText: 'Too Many Requests',
340-
text: async () => 'Rate limit exceeded',
341-
} as any)
325+
text: 'Rate limit exceeded',
326+
})
342327

343328
await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
344329

@@ -356,17 +341,16 @@ describe('Knowledge Search Utils', () => {
356341
KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
357342
})
358343

359-
const fetchSpy = vi.mocked(fetch)
360-
fetchSpy.mockResolvedValueOnce({
361-
ok: true,
362-
json: async () => ({
344+
mockNextFetchResponse({
345+
json: {
363346
data: [{ embedding: [0.1, 0.2, 0.3] }],
364-
}),
365-
} as any)
347+
usage: { prompt_tokens: 1, total_tokens: 1 },
348+
},
349+
})
366350

367351
await generateSearchEmbedding('test query')
368352

369-
expect(fetchSpy).toHaveBeenCalledWith(
353+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
370354
expect.any(String),
371355
expect.objectContaining({
372356
body: JSON.stringify({
@@ -387,17 +371,16 @@ describe('Knowledge Search Utils', () => {
387371
OPENAI_API_KEY: 'test-openai-key',
388372
})
389373

390-
const fetchSpy = vi.mocked(fetch)
391-
fetchSpy.mockResolvedValueOnce({
392-
ok: true,
393-
json: async () => ({
374+
mockNextFetchResponse({
375+
json: {
394376
data: [{ embedding: [0.1, 0.2, 0.3] }],
395-
}),
396-
} as any)
377+
usage: { prompt_tokens: 1, total_tokens: 1 },
378+
},
379+
})
397380

398381
await generateSearchEmbedding('test query', 'text-embedding-3-small')
399382

400-
expect(fetchSpy).toHaveBeenCalledWith(
383+
expect(vi.mocked(fetch)).toHaveBeenCalledWith(
401384
expect.any(String),
402385
expect.objectContaining({
403386
body: JSON.stringify({

‎apps/sim/lib/billing/core/subscription.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -448,9 +448,11 @@ export async function hasInboxAccess(userId: string): Promise<boolean> {
448448
if (!isProd) {
449449
return true
450450
}
451-
const sub = await getHighestPrioritySubscription(userId)
451+
const [sub, billingStatus] = await Promise.all([
452+
getHighestPrioritySubscription(userId),
453+
getEffectiveBillingStatus(userId),
454+
])
452455
if (!sub) return false
453-
const billingStatus = await getEffectiveBillingStatus(userId)
454456
if (!hasUsableSubscriptionAccess(sub.status, billingStatus.billingBlocked)) return false
455457
return getPlanTierCredits(sub.plan) >= 25000 || checkEnterprisePlan(sub)
456458
} catch (error) {
@@ -470,9 +472,11 @@ export async function hasLiveSyncAccess(userId: string): Promise<boolean> {
470472
if (!isHosted) {
471473
return true
472474
}
473-
const sub = await getHighestPrioritySubscription(userId)
475+
const [sub, billingStatus] = await Promise.all([
476+
getHighestPrioritySubscription(userId),
477+
getEffectiveBillingStatus(userId),
478+
])
474479
if (!sub) return false
475-
const billingStatus = await getEffectiveBillingStatus(userId)
476480
if (!hasUsableSubscriptionAccess(sub.status, billingStatus.billingBlocked)) return false
477481
return getPlanTierCredits(sub.plan) >= 25000 || checkEnterprisePlan(sub)
478482
} catch (error) {

0 commit comments

Comments
 (0)