wehub-resource-sync
365 lines
13 KiB
Swift
365 lines
13 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
import Testing
|
|
|
|
struct LiteLLMUsageFetcherTests {
|
|
@Test
|
|
func `parses user usage with personal and team budgets`() throws {
|
|
let json = """
|
|
{
|
|
"user_id": "user-123",
|
|
"user_info": {
|
|
"user_id": "user-123",
|
|
"user_alias": "litellm-user@example.com",
|
|
"max_budget": 300.0,
|
|
"spend": 212.3537162499998,
|
|
"user_email": "litellm-user@example.com",
|
|
"budget_reset_at": null,
|
|
"teams": ["team-456"],
|
|
"metadata": {
|
|
"source": "keycloak",
|
|
"preferred_username": "litellm-user@example.com",
|
|
"budget": 300,
|
|
"flags": {
|
|
"keycloak": true
|
|
}
|
|
}
|
|
},
|
|
"keys": [
|
|
{
|
|
"key_name": "sk-...OTHER",
|
|
"user_id": "user-123",
|
|
"team_id": "team-other"
|
|
},
|
|
{
|
|
"key_name": "sk-...IAAw",
|
|
"spend": 212.3537162499998,
|
|
"expires": "2026-09-11T00:12:55.950000+00:00",
|
|
"user_id": "user-123",
|
|
"team_id": "team-456"
|
|
}
|
|
],
|
|
"teams": [
|
|
{
|
|
"team_alias": "unrelated",
|
|
"team_id": "team-other",
|
|
"max_budget": 5.0,
|
|
"spend": 4.0
|
|
},
|
|
{
|
|
"team_alias": "ai",
|
|
"team_id": "team-456",
|
|
"max_budget": 1000.0,
|
|
"spend": 215.3245658499998,
|
|
"budget_duration": "7d",
|
|
"budget_reset_at": "2026-06-15T00:00:00Z"
|
|
}
|
|
]
|
|
}
|
|
"""
|
|
|
|
let parsed = try LiteLLMUsageFetcher._parseUserInfoForTesting(
|
|
Data(json.utf8),
|
|
keyInfo: LiteLLMKeyInfoSnapshot(
|
|
userID: "user-123",
|
|
teamID: "team-456",
|
|
keyName: "sk-...IAAw",
|
|
spendUSD: 212.3537162499998,
|
|
expiresAt: Date(timeIntervalSince1970: 2)),
|
|
updatedAt: Date(timeIntervalSince1970: 1))
|
|
|
|
#expect(parsed.userID == "user-123")
|
|
#expect(parsed.accountEmail == "litellm-user@example.com")
|
|
#expect(abs(parsed.personalSpendUSD - 212.3537162499998) < 0.000001)
|
|
#expect(parsed.personalBudgetUSD == 300)
|
|
#expect(parsed.teamUsage?.alias == "ai")
|
|
#expect(parsed.teamUsage?.spendUSD == 215.3245658499998)
|
|
#expect(parsed.teamUsage?.budgetUSD == 1000)
|
|
#expect(parsed.keyName == "sk-...IAAw")
|
|
#expect(parsed.keyExpiresAt == Date(timeIntervalSince1970: 2))
|
|
|
|
let snapshot = parsed.toUsageSnapshot()
|
|
#expect(snapshot.identity?.providerID == .litellm)
|
|
#expect(snapshot.identity?.accountEmail == "litellm-user@example.com")
|
|
let primary = try #require(snapshot.primary)
|
|
#expect(abs(primary.usedPercent - 70.78457208333327) < 0.000001)
|
|
#expect(primary.resetDescription == "$212.35 / $300.00")
|
|
let secondary = try #require(snapshot.secondary)
|
|
#expect(abs(secondary.usedPercent - 21.53245658499998) < 0.000001)
|
|
#expect(secondary.resetDescription == "Team ai: $215.32 / $1,000.00")
|
|
#expect(snapshot.providerCost?.used == 212.3537162499998)
|
|
#expect(snapshot.providerCost?.limit == 300)
|
|
#expect(snapshot.providerCost?.period == "Personal budget")
|
|
}
|
|
|
|
@Test
|
|
func `preserves personal spend when no budget is configured`() throws {
|
|
let json = """
|
|
{
|
|
"user_id": "user-123",
|
|
"user_info": {
|
|
"user_id": "user-123",
|
|
"max_budget": null,
|
|
"spend": 12.5
|
|
}
|
|
}
|
|
"""
|
|
|
|
let parsed = try LiteLLMUsageFetcher._parseUserInfoForTesting(
|
|
Data(json.utf8),
|
|
keyInfo: LiteLLMKeyInfoSnapshot(
|
|
userID: "user-123",
|
|
teamID: nil,
|
|
keyName: "personal-key",
|
|
spendUSD: 12.5,
|
|
expiresAt: nil),
|
|
updatedAt: Date(timeIntervalSince1970: 1))
|
|
|
|
let snapshot = parsed.toUsageSnapshot()
|
|
#expect(snapshot.primary == nil)
|
|
#expect(snapshot.providerCost?.used == 12.5)
|
|
#expect(snapshot.providerCost?.limit == 0)
|
|
#expect(snapshot.providerCost?.period == "Personal spend")
|
|
}
|
|
|
|
@Test
|
|
func `parses key info identity for user lookup`() throws {
|
|
let json = """
|
|
{
|
|
"key": "sk-redacted",
|
|
"info": {
|
|
"key_name": "sk-...IAAw",
|
|
"spend": 212.3537162499998,
|
|
"expires": "2026-09-11T00:12:55.950000+00:00",
|
|
"user_id": "user-123",
|
|
"team_id": "team-456",
|
|
"max_budget": null
|
|
}
|
|
}
|
|
"""
|
|
|
|
let parsed = try LiteLLMUsageFetcher._parseKeyInfoForTesting(Data(json.utf8))
|
|
|
|
#expect(parsed.userID == "user-123")
|
|
#expect(parsed.teamID == "team-456")
|
|
#expect(parsed.keyName == "sk-...IAAw")
|
|
#expect(parsed.spendUSD == 212.3537162499998)
|
|
}
|
|
|
|
@Test
|
|
func `parses team-only key info without user identity`() throws {
|
|
let json = """
|
|
{
|
|
"info": {
|
|
"key_name": "team-service-key",
|
|
"spend": 25.0,
|
|
"team_id": "team-456"
|
|
}
|
|
}
|
|
"""
|
|
|
|
let parsed = try LiteLLMUsageFetcher._parseKeyInfoForTesting(Data(json.utf8))
|
|
|
|
#expect(parsed.userID == nil)
|
|
#expect(parsed.teamID == "team-456")
|
|
#expect(parsed.keyName == "team-service-key")
|
|
}
|
|
|
|
@Test
|
|
func `management urls accept root or v1 base urls`() throws {
|
|
let root = try #require(URL(string: "https://litellm.example.com"))
|
|
let versioned = try #require(URL(string: "https://litellm.example.com/v1"))
|
|
let nestedVersioned = try #require(URL(string: "https://gateway.example.com/litellm/v1/"))
|
|
|
|
#expect(
|
|
LiteLLMUsageFetcher
|
|
._keyInfoURLForTesting(baseURL: root)
|
|
.absoluteString == "https://litellm.example.com/key/info")
|
|
#expect(
|
|
LiteLLMUsageFetcher
|
|
._keyInfoURLForTesting(baseURL: versioned)
|
|
.absoluteString == "https://litellm.example.com/key/info")
|
|
#expect(
|
|
LiteLLMUsageFetcher
|
|
._userInfoURLForTesting(baseURL: nestedVersioned, userID: "user-123")
|
|
.absoluteString == "https://gateway.example.com/litellm/user/info?user_id=user-123")
|
|
#expect(
|
|
LiteLLMUsageFetcher
|
|
._teamInfoURLForTesting(baseURL: nestedVersioned, teamID: "team-456")
|
|
.absoluteString == "https://gateway.example.com/litellm/team/info?team_id=team-456")
|
|
}
|
|
|
|
@Test
|
|
func `settings reader trims quoted environment values`() {
|
|
let environment = [
|
|
"LITELLM_API_KEY": " 'sk-test' ",
|
|
"LITELLM_BASE_URL": #" "https://litellm.example.com/v1" "#,
|
|
]
|
|
|
|
#expect(LiteLLMSettingsReader.apiKey(environment: environment) == "sk-test")
|
|
#expect(LiteLLMSettingsReader.baseURL(environment: environment)?
|
|
.absoluteString == "https://litellm.example.com/v1")
|
|
}
|
|
|
|
@Test
|
|
func `fetch trims api key before sending management requests`() async throws {
|
|
let baseURL = try #require(URL(string: "https://litellm.example.com/v1"))
|
|
let transport = ProviderHTTPTransportStub { request in
|
|
#expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer sk-test")
|
|
|
|
let path = request.url?.path
|
|
let query = request.url?.query
|
|
let body: String
|
|
switch path {
|
|
case "/key/info":
|
|
#expect(query == nil)
|
|
body = """
|
|
{
|
|
"info": {
|
|
"user_id": "user-123",
|
|
"team_id": "team-456",
|
|
"spend": 1
|
|
}
|
|
}
|
|
"""
|
|
case "/user/info":
|
|
#expect(query == "user_id=user-123")
|
|
body = """
|
|
{
|
|
"user_id": "user-123",
|
|
"user_info": {
|
|
"user_id": "user-123",
|
|
"max_budget": 10,
|
|
"spend": 1
|
|
}
|
|
}
|
|
"""
|
|
default:
|
|
Issue.record("unexpected LiteLLM request path: \(path ?? "nil")")
|
|
body = "{}"
|
|
}
|
|
|
|
let url = try #require(request.url)
|
|
let response = try #require(HTTPURLResponse(
|
|
url: url,
|
|
statusCode: 200,
|
|
httpVersion: nil,
|
|
headerFields: nil))
|
|
return (Data(body.utf8), response)
|
|
}
|
|
|
|
let snapshot = try await LiteLLMUsageFetcher.fetchUsage(
|
|
apiKey: " sk-test\n",
|
|
baseURL: baseURL,
|
|
transport: transport)
|
|
|
|
#expect(snapshot.userID == "user-123")
|
|
let requests = await transport.requests()
|
|
#expect(requests.count == 2)
|
|
}
|
|
|
|
@Test
|
|
func `fetches team usage for team-only virtual keys`() async throws {
|
|
let baseURL = try #require(URL(string: "https://litellm.example.com/v1"))
|
|
let transport = ProviderHTTPTransportStub { request in
|
|
#expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer sk-team")
|
|
|
|
let path = request.url?.path
|
|
let query = request.url?.query
|
|
let body: String
|
|
switch path {
|
|
case "/key/info":
|
|
#expect(query == nil)
|
|
body = """
|
|
{
|
|
"info": {
|
|
"key_name": "team-service-key",
|
|
"team_id": "team-456",
|
|
"spend": 25
|
|
}
|
|
}
|
|
"""
|
|
case "/team/info":
|
|
#expect(query == "team_id=team-456")
|
|
body = """
|
|
{
|
|
"team_id": "team-456",
|
|
"team_info": {
|
|
"team_id": "team-456",
|
|
"team_alias": "platform",
|
|
"max_budget": 100,
|
|
"spend": 25,
|
|
"budget_duration": "30d",
|
|
"budget_reset_at": "2026-07-01T00:00:00Z"
|
|
}
|
|
}
|
|
"""
|
|
default:
|
|
Issue.record("unexpected LiteLLM request path: \(path ?? "nil")")
|
|
body = "{}"
|
|
}
|
|
|
|
let url = try #require(request.url)
|
|
let response = try #require(HTTPURLResponse(
|
|
url: url,
|
|
statusCode: 200,
|
|
httpVersion: nil,
|
|
headerFields: nil))
|
|
return (Data(body.utf8), response)
|
|
}
|
|
|
|
let snapshot = try await LiteLLMUsageFetcher.fetchUsage(
|
|
apiKey: "sk-team",
|
|
baseURL: baseURL,
|
|
transport: transport,
|
|
updatedAt: Date(timeIntervalSince1970: 1))
|
|
|
|
#expect(snapshot.userID == nil)
|
|
#expect(snapshot.teamUsage?.id == "team-456")
|
|
#expect(snapshot.teamUsage?.alias == "platform")
|
|
#expect(snapshot.teamUsage?.spendUSD == 25)
|
|
#expect(snapshot.teamUsage?.budgetUSD == 100)
|
|
|
|
let usage = snapshot.toUsageSnapshot()
|
|
#expect(usage.primary == nil)
|
|
#expect(usage.secondary?.usedPercent == 25)
|
|
#expect(usage.providerCost?.used == 25)
|
|
#expect(usage.providerCost?.limit == 100)
|
|
#expect(usage.providerCost?.period == "Team budget")
|
|
|
|
let requests = await transport.requests()
|
|
#expect(requests.count == 2)
|
|
}
|
|
|
|
@Test
|
|
func `fetch surfaces rejected virtual key`() async throws {
|
|
let baseURL = try #require(URL(string: "https://litellm.example.com"))
|
|
let transport = ProviderHTTPTransportStub { request in
|
|
#expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer sk-target")
|
|
let url = try #require(request.url)
|
|
let response = try #require(HTTPURLResponse(
|
|
url: url,
|
|
statusCode: 401,
|
|
httpVersion: nil,
|
|
headerFields: nil))
|
|
return (Data(#"{"detail":"Unauthorized"}"#.utf8), response)
|
|
}
|
|
|
|
do {
|
|
_ = try await LiteLLMUsageFetcher.fetchUsage(
|
|
apiKey: "sk-target",
|
|
baseURL: baseURL,
|
|
transport: transport)
|
|
Issue.record("expected LiteLLMUsageError.apiError")
|
|
} catch let LiteLLMUsageError.apiError(message) {
|
|
#expect(message.contains("HTTP 401"))
|
|
#expect(message.contains("Unauthorized"))
|
|
} catch {
|
|
Issue.record("expected LiteLLMUsageError.apiError, got \(error)")
|
|
}
|
|
|
|
let requests = await transport.requests()
|
|
#expect(requests.count == 1)
|
|
}
|
|
}
|