97 lines
2.5 KiB
TypeScript
97 lines
2.5 KiB
TypeScript
import fs from 'node:fs'
|
|
import os from 'node:os'
|
|
import path from 'node:path'
|
|
import initSqlite from 'sql.js'
|
|
|
|
export const API_BASE = `http://localhost:${process.env.NEX_BACKEND_PORT || '19026'}`
|
|
|
|
export interface SeedProviderInput {
|
|
id: string
|
|
name: string
|
|
apiKey: string
|
|
baseUrl: string
|
|
protocol: 'openai' | 'anthropic'
|
|
enabled: boolean
|
|
}
|
|
|
|
export interface SeedModelInput {
|
|
providerId: string
|
|
modelName: string
|
|
enabled: boolean
|
|
}
|
|
|
|
export interface SeedStatsInput {
|
|
providerId: string
|
|
modelName: string
|
|
requestCount: number
|
|
date: string
|
|
}
|
|
|
|
export async function clearDatabase(request: import('@playwright/test').APIRequestContext) {
|
|
const providers = await request.get(`${API_BASE}/api/providers`)
|
|
if (providers.ok()) {
|
|
const data = await providers.json()
|
|
for (const p of data) {
|
|
await request.delete(`${API_BASE}/api/providers/${p.id}`)
|
|
}
|
|
}
|
|
}
|
|
|
|
export async function seedProvider(request: import('@playwright/test').APIRequestContext, data: SeedProviderInput) {
|
|
const resp = await request.post(`${API_BASE}/api/providers`, {
|
|
data: {
|
|
id: data.id,
|
|
name: data.name,
|
|
api_key: data.apiKey,
|
|
base_url: data.baseUrl,
|
|
protocol: data.protocol,
|
|
enabled: data.enabled,
|
|
},
|
|
})
|
|
if (!resp.ok()) {
|
|
throw new Error(`seedProvider failed: ${resp.status()} ${await resp.text()}`)
|
|
}
|
|
return resp.json()
|
|
}
|
|
|
|
export async function seedModel(request: import('@playwright/test').APIRequestContext, data: SeedModelInput) {
|
|
const resp = await request.post(`${API_BASE}/api/models`, {
|
|
data: {
|
|
provider_id: data.providerId,
|
|
model_name: data.modelName,
|
|
enabled: data.enabled,
|
|
},
|
|
})
|
|
if (!resp.ok()) {
|
|
throw new Error(`seedModel failed: ${resp.status()} ${await resp.text()}`)
|
|
}
|
|
return resp.json()
|
|
}
|
|
|
|
export async function seedUsageStats(statsData: SeedStatsInput[]) {
|
|
const tempDir = path.join(os.tmpdir(), 'nex-e2e')
|
|
|
|
const dbPath = path.join(tempDir, 'test.db')
|
|
|
|
if (!fs.existsSync(dbPath)) {
|
|
throw new Error(`Database file not found at ${dbPath}. Backend may not have created it yet.`)
|
|
}
|
|
|
|
const SQL = await initSqlite()
|
|
const buf = fs.readFileSync(dbPath)
|
|
const db = new SQL.Database(buf)
|
|
|
|
for (const row of statsData) {
|
|
db.run('INSERT OR REPLACE INTO usage_stats (provider_id, model_name, request_count, date) VALUES (?, ?, ?, ?)', [
|
|
row.providerId,
|
|
row.modelName,
|
|
row.requestCount,
|
|
row.date,
|
|
])
|
|
}
|
|
|
|
const data = db.export()
|
|
fs.writeFileSync(dbPath, Buffer.from(data))
|
|
db.close()
|
|
}
|