1693 lines
47 KiB
Go
1693 lines
47 KiB
Go
package governance
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// ============================================================================
|
|
// SCENARIO 1: VK Switching Teams After Budget Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestVKSwitchTeamAfterBudgetExhaustion verifies that after exhausting one team's budget,
|
|
// switching the VK to another team allows requests to pass
|
|
func TestVKSwitchTeamAfterBudgetExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create Team 1 with small budget
|
|
team1Name := "test-team1-switch-" + generateRandomID()
|
|
team1Budget := 0.01 // $0.01
|
|
createTeam1Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: team1Name,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: team1Budget,
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeam1Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team1: status %d", createTeam1Resp.StatusCode)
|
|
}
|
|
|
|
team1ID := ExtractIDFromResponse(t, createTeam1Resp)
|
|
testData.AddTeam(team1ID)
|
|
|
|
// Create Team 2 with higher budget
|
|
team2Name := "test-team2-switch-" + generateRandomID()
|
|
team2Budget := 10.0 // $10
|
|
createTeam2Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: team2Name,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: team2Budget,
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeam2Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team2: status %d", createTeam2Resp.StatusCode)
|
|
}
|
|
|
|
team2ID := ExtractIDFromResponse(t, createTeam2Resp)
|
|
testData.AddTeam(team2ID)
|
|
|
|
t.Logf("Created Team1 (budget: $%.2f) and Team2 (budget: $%.2f)", team1Budget, team2Budget)
|
|
|
|
// Create VK assigned to Team 1
|
|
vkName := "test-vk-team-switch-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
TeamID: &team1ID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created VK assigned to Team1")
|
|
|
|
// Exhaust Team1's budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
t.Logf("Team1 budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
|
|
if consumedBudget >= team1Budget {
|
|
// Make one more request to trigger rejection
|
|
continue
|
|
}
|
|
}
|
|
|
|
if consumedBudget < team1Budget {
|
|
t.Fatalf("Could not exhaust Team1 budget")
|
|
}
|
|
|
|
// Now switch VK to Team2
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
Body: UpdateVirtualKeyRequest{
|
|
TeamID: &team2ID,
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to switch VK to Team2: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Switched VK from Team1 to Team2")
|
|
|
|
// Wait for in-memory update
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed with Team2's budget
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after switching to Team2"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after switching to Team2 with available budget, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("VK switch team after budget exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 2: VK Switching Customers After Budget Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestVKSwitchCustomerAfterBudgetExhaustion verifies that after exhausting one customer's budget,
|
|
// switching the VK to another customer allows requests to pass
|
|
func TestVKSwitchCustomerAfterBudgetExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create Customer 1 with small budget
|
|
customer1Name := "test-customer1-switch-" + generateRandomID()
|
|
customer1Budget := 0.01 // $0.01
|
|
createCustomer1Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customer1Name,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: customer1Budget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomer1Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer1: status %d", createCustomer1Resp.StatusCode)
|
|
}
|
|
|
|
customer1ID := ExtractIDFromResponse(t, createCustomer1Resp)
|
|
testData.AddCustomer(customer1ID)
|
|
|
|
// Create Customer 2 with higher budget
|
|
customer2Name := "test-customer2-switch-" + generateRandomID()
|
|
customer2Budget := 10.0 // $10
|
|
createCustomer2Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customer2Name,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: customer2Budget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomer2Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer2: status %d", createCustomer2Resp.StatusCode)
|
|
}
|
|
|
|
customer2ID := ExtractIDFromResponse(t, createCustomer2Resp)
|
|
testData.AddCustomer(customer2ID)
|
|
|
|
t.Logf("Created Customer1 (budget: $%.2f) and Customer2 (budget: $%.2f)", customer1Budget, customer2Budget)
|
|
|
|
// Create VK assigned directly to Customer 1
|
|
vkName := "test-vk-customer-switch-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
CustomerID: &customer1ID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created VK assigned to Customer1")
|
|
|
|
// Exhaust Customer1's budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
t.Logf("Customer1 budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
|
|
if consumedBudget >= customer1Budget {
|
|
continue
|
|
}
|
|
}
|
|
|
|
if consumedBudget < customer1Budget {
|
|
t.Fatalf("Could not exhaust Customer1 budget")
|
|
}
|
|
|
|
// Now switch VK to Customer2
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
Body: UpdateVirtualKeyRequest{
|
|
CustomerID: &customer2ID,
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to switch VK to Customer2: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Switched VK from Customer1 to Customer2")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed with Customer2's budget
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after switching to Customer2"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after switching to Customer2 with available budget, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("VK switch customer after budget exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 3: Hierarchical Chain VK->Team->Customer Budget Switching
|
|
// ============================================================================
|
|
|
|
// TestHierarchicalChainBudgetSwitch verifies switching the entire hierarchy
|
|
func TestHierarchicalChainBudgetSwitch(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create Customer 1 with small budget
|
|
customer1Name := "test-customer1-hierarchy-" + generateRandomID()
|
|
createCustomer1Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customer1Name,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: 0.01, // $0.01 - most restrictive
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomer1Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer1: status %d", createCustomer1Resp.StatusCode)
|
|
}
|
|
|
|
customer1ID := ExtractIDFromResponse(t, createCustomer1Resp)
|
|
testData.AddCustomer(customer1ID)
|
|
|
|
// Create Team 1 under Customer 1
|
|
team1Name := "test-team1-hierarchy-" + generateRandomID()
|
|
createTeam1Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: team1Name,
|
|
CustomerID: &customer1ID,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: 100.0, // High budget - customer is limiting
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeam1Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team1: status %d", createTeam1Resp.StatusCode)
|
|
}
|
|
|
|
team1ID := ExtractIDFromResponse(t, createTeam1Resp)
|
|
testData.AddTeam(team1ID)
|
|
|
|
// Create Customer 2 with higher budget
|
|
customer2Name := "test-customer2-hierarchy-" + generateRandomID()
|
|
createCustomer2Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customer2Name,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: 100.0, // High budget
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomer2Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer2: status %d", createCustomer2Resp.StatusCode)
|
|
}
|
|
|
|
customer2ID := ExtractIDFromResponse(t, createCustomer2Resp)
|
|
testData.AddCustomer(customer2ID)
|
|
|
|
// Create Team 2 under Customer 2
|
|
team2Name := "test-team2-hierarchy-" + generateRandomID()
|
|
createTeam2Resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: team2Name,
|
|
CustomerID: &customer2ID,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: 100.0, // High budget
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeam2Resp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team2: status %d", createTeam2Resp.StatusCode)
|
|
}
|
|
|
|
team2ID := ExtractIDFromResponse(t, createTeam2Resp)
|
|
testData.AddTeam(team2ID)
|
|
|
|
t.Logf("Created hierarchy: Customer1(low budget)->Team1 and Customer2(high budget)->Team2")
|
|
|
|
// Create VK assigned to Team 1
|
|
vkName := "test-vk-hierarchy-switch-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
TeamID: &team1ID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
// Exhaust Customer1's budget (which is limiting Team1)
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
budgetExhausted := false
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
budgetExhausted = true
|
|
t.Logf("Customer1 budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
}
|
|
|
|
if !budgetExhausted {
|
|
t.Fatalf("Budget should have been exhausted within 150 requests, but no budget rejection was observed (consumed: $%.6f)", consumedBudget)
|
|
}
|
|
|
|
// Switch VK to Team2 (under Customer2)
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
Body: UpdateVirtualKeyRequest{
|
|
TeamID: &team2ID,
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to switch VK to Team2: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Switched VK from Team1(Customer1) to Team2(Customer2)")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after switching hierarchy"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after switching hierarchy, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Hierarchical chain budget switch verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 4: VK Budget Update After Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestVKBudgetUpdateAfterExhaustion verifies that updating VK budget after exhaustion allows requests
|
|
func TestVKBudgetUpdateAfterExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create VK with small budget
|
|
vkName := "test-vk-budget-update-" + generateRandomID()
|
|
initialBudget := 0.01 // $0.01
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: initialBudget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created VK with budget: $%.2f", initialBudget)
|
|
|
|
// Exhaust VK budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
sawBudgetRejection := false
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
sawBudgetRejection = true
|
|
t.Logf("VK budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
}
|
|
|
|
if !sawBudgetRejection {
|
|
t.Fatalf("No budget rejection observed; consumed budget: $%.6f", consumedBudget)
|
|
}
|
|
|
|
// Update VK budget to a higher value
|
|
newBudget := 10.0
|
|
resetDuration := "1h"
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
Body: UpdateVirtualKeyRequest{
|
|
Budget: &UpdateBudgetRequest{
|
|
MaxLimit: &newBudget,
|
|
ResetDuration: &resetDuration,
|
|
},
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to update VK budget: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Updated VK budget from $%.2f to $%.2f", initialBudget, newBudget)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after budget update"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after budget update, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("VK budget update after exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 5: Team Budget Update After Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestTeamBudgetUpdateAfterExhaustion verifies that updating team budget after exhaustion allows requests
|
|
func TestTeamBudgetUpdateAfterExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create team with small budget
|
|
teamName := "test-team-budget-update-" + generateRandomID()
|
|
initialBudget := 0.01 // $0.01
|
|
createTeamResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: teamName,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: initialBudget,
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeamResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team: status %d", createTeamResp.StatusCode)
|
|
}
|
|
|
|
teamID := ExtractIDFromResponse(t, createTeamResp)
|
|
testData.AddTeam(teamID)
|
|
|
|
// Create VK under team
|
|
vkName := "test-vk-team-budget-update-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
TeamID: &teamID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created team with budget: $%.2f", initialBudget)
|
|
|
|
// Exhaust team budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
sawBudgetRejection := false
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
sawBudgetRejection = true
|
|
t.Logf("Team budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
}
|
|
|
|
if !sawBudgetRejection {
|
|
t.Fatalf("No budget rejection observed; consumed budget: $%.6f", consumedBudget)
|
|
}
|
|
|
|
// Update team budget
|
|
newBudget := 10.0
|
|
resetDuration := "1h"
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/teams/" + teamID,
|
|
Body: UpdateTeamRequest{
|
|
Budgets: &[]BudgetRequest{{
|
|
MaxLimit: newBudget,
|
|
ResetDuration: resetDuration,
|
|
}},
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to update team budget: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Updated team budget from $%.2f to $%.2f", initialBudget, newBudget)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after team budget update"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after team budget update, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Team budget update after exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 6: Customer Budget Update After Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestCustomerBudgetUpdateAfterExhaustion verifies that updating customer budget after exhaustion allows requests
|
|
func TestCustomerBudgetUpdateAfterExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create customer with small budget
|
|
customerName := "test-customer-budget-update-" + generateRandomID()
|
|
initialBudget := 0.01 // $0.01
|
|
createCustomerResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customerName,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: initialBudget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomerResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer: status %d", createCustomerResp.StatusCode)
|
|
}
|
|
|
|
customerID := ExtractIDFromResponse(t, createCustomerResp)
|
|
testData.AddCustomer(customerID)
|
|
|
|
// Create team under customer
|
|
teamName := "test-team-customer-update-" + generateRandomID()
|
|
createTeamResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: teamName,
|
|
CustomerID: &customerID,
|
|
},
|
|
})
|
|
|
|
if createTeamResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team: status %d", createTeamResp.StatusCode)
|
|
}
|
|
|
|
teamID := ExtractIDFromResponse(t, createTeamResp)
|
|
testData.AddTeam(teamID)
|
|
|
|
// Create VK under team
|
|
vkName := "test-vk-customer-budget-update-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
TeamID: &teamID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created customer with budget: $%.2f", initialBudget)
|
|
|
|
// Exhaust customer budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
sawBudgetRejection := false
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
sawBudgetRejection = true
|
|
t.Logf("Customer budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
}
|
|
|
|
if !sawBudgetRejection {
|
|
t.Fatalf("No budget rejection observed; consumed budget: $%.6f", consumedBudget)
|
|
}
|
|
|
|
// Update customer budget
|
|
newBudget := 10.0
|
|
resetDuration := "1h"
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/customers/" + customerID,
|
|
Body: UpdateCustomerRequest{
|
|
Budget: &UpdateBudgetRequest{
|
|
MaxLimit: &newBudget,
|
|
ResetDuration: &resetDuration,
|
|
},
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to update customer budget: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Updated customer budget from $%.2f to $%.2f", initialBudget, newBudget)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after customer budget update"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after customer budget update, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Customer budget update after exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 7: Provider Config Budget Update After Exhaustion
|
|
// ============================================================================
|
|
|
|
// TestProviderConfigBudgetUpdateAfterExhaustion verifies that updating provider config budget after exhaustion allows requests
|
|
func TestProviderConfigBudgetUpdateAfterExhaustion(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create VK with provider config budget
|
|
vkName := "test-vk-provider-budget-update-" + generateRandomID()
|
|
initialBudget := 0.01 // $0.01
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
ProviderConfigs: []ProviderConfigRequest{
|
|
{
|
|
Provider: "openai",
|
|
Weight: float64Ptr(1.0),
|
|
AllowedModels: []string{"*"},
|
|
KeyIDs: []string{"*"},
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: initialBudget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created VK with provider config budget: $%.2f", initialBudget)
|
|
|
|
// Get provider config ID
|
|
getDataResp := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap := getDataResp.Body["virtual_keys"].(map[string]interface{})
|
|
vkData := virtualKeysMap[vkValue].(map[string]interface{})
|
|
providerConfigs := vkData["provider_configs"].([]interface{})
|
|
providerConfig := providerConfigs[0].(map[string]interface{})
|
|
providerConfigID := uint(providerConfig["id"].(float64))
|
|
|
|
// Exhaust provider config budget
|
|
consumedBudget := 0.0
|
|
requestNum := 1
|
|
sawBudgetRejection := false
|
|
|
|
for requestNum <= 150 {
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Hello how are you?"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode >= 400 {
|
|
if CheckErrorMessage(t, resp, "budget") {
|
|
sawBudgetRejection = true
|
|
t.Logf("Provider config budget exhausted at request %d (consumed: $%.6f)", requestNum, consumedBudget)
|
|
break
|
|
} else {
|
|
t.Fatalf("Request %d failed with unexpected error: %v", requestNum, resp.Body)
|
|
}
|
|
}
|
|
|
|
if usage, ok := resp.Body["usage"].(map[string]interface{}); ok {
|
|
if prompt, ok := usage["prompt_tokens"].(float64); ok {
|
|
if completion, ok := usage["completion_tokens"].(float64); ok {
|
|
cost, _ := CalculateCost("openai/gpt-4o", int(prompt), int(completion))
|
|
consumedBudget += cost
|
|
}
|
|
}
|
|
}
|
|
|
|
requestNum++
|
|
}
|
|
|
|
if !sawBudgetRejection {
|
|
t.Fatalf("No budget rejection observed; consumed budget: $%.6f", consumedBudget)
|
|
}
|
|
|
|
// Update provider config budget
|
|
newBudget := 10.0
|
|
updateResp := MakeRequest(t, APIRequest{
|
|
Method: "PUT",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
Body: UpdateVirtualKeyRequest{
|
|
ProviderConfigs: []ProviderConfigRequest{
|
|
{
|
|
ID: &providerConfigID,
|
|
Provider: "openai",
|
|
Weight: float64Ptr(1.0),
|
|
AllowedModels: []string{"*"},
|
|
KeyIDs: []string{"*"},
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: newBudget,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
if updateResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to update provider config budget: status %d", updateResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Updated provider config budget from $%.2f to $%.2f", initialBudget, newBudget)
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Request should now succeed
|
|
resp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/v1/chat/completions",
|
|
Body: ChatCompletionRequest{
|
|
Model: "openai/gpt-4o",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: "Request after provider config budget update"},
|
|
},
|
|
},
|
|
VKHeader: &vkValue,
|
|
})
|
|
|
|
if resp.StatusCode != 200 {
|
|
t.Fatalf("Request should succeed after provider config budget update, got status %d", resp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Provider config budget update after exhaustion verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 8: VK Deletion Cascade
|
|
// ============================================================================
|
|
|
|
// TestVKDeletionCascadeComplete verifies deleting VK removes provider configs, budgets, and rate limits from memory
|
|
func TestVKDeletionCascadeComplete(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create VK with budget, rate limit, and provider configs
|
|
vkName := "test-vk-deletion-cascade-" + generateRandomID()
|
|
tokenLimit := int64(10000)
|
|
tokenResetDuration := "1h"
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: 10.0,
|
|
ResetDuration: "1h",
|
|
},
|
|
RateLimit: &CreateRateLimitRequest{
|
|
TokenMaxLimit: &tokenLimit,
|
|
TokenResetDuration: &tokenResetDuration,
|
|
},
|
|
ProviderConfigs: []ProviderConfigRequest{
|
|
{
|
|
Provider: "openai",
|
|
Weight: float64Ptr(1.0),
|
|
AllowedModels: []string{"*"},
|
|
KeyIDs: []string{"*"},
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: 5.0,
|
|
ResetDuration: "1h",
|
|
},
|
|
RateLimit: &CreateRateLimitRequest{
|
|
TokenMaxLimit: &tokenLimit,
|
|
TokenResetDuration: &tokenResetDuration,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
// Don't add to testData since we'll delete manually
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created VK with budget, rate limit, and provider config")
|
|
|
|
// Get initial state from in-memory store
|
|
getDataResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap1 := getDataResp1.Body["virtual_keys"].(map[string]interface{})
|
|
|
|
getBudgetsResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
budgetsMap1 := getBudgetsResp1.Body["budgets"].(map[string]interface{})
|
|
|
|
getRateLimitsResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/rate-limits?from_memory=true",
|
|
})
|
|
|
|
rateLimitsMap1 := getRateLimitsResp1.Body["rate_limits"].(map[string]interface{})
|
|
|
|
// Verify VK exists
|
|
_, vkExists := virtualKeysMap1[vkValue]
|
|
if !vkExists {
|
|
t.Fatalf("VK not found in in-memory store")
|
|
}
|
|
|
|
vkData1 := virtualKeysMap1[vkValue].(map[string]interface{})
|
|
vkBudgetID := vkData1["budget_id"].(string)
|
|
vkRateLimitID := vkData1["rate_limit_id"].(string)
|
|
providerConfigs := vkData1["provider_configs"].([]interface{})
|
|
pc := providerConfigs[0].(map[string]interface{})
|
|
pcBudgetID := pc["budget_id"].(string)
|
|
pcRateLimitID := pc["rate_limit_id"].(string)
|
|
|
|
// Verify all resources exist in memory
|
|
_, vkBudgetExists := budgetsMap1[vkBudgetID]
|
|
_, vkRateLimitExists := rateLimitsMap1[vkRateLimitID]
|
|
_, pcBudgetExists := budgetsMap1[pcBudgetID]
|
|
_, pcRateLimitExists := rateLimitsMap1[pcRateLimitID]
|
|
|
|
if !vkBudgetExists || !vkRateLimitExists || !pcBudgetExists || !pcRateLimitExists {
|
|
t.Fatalf("Not all resources found in memory before deletion")
|
|
}
|
|
|
|
t.Logf("All resources exist in memory before deletion ✓")
|
|
|
|
// Delete VK
|
|
deleteResp := MakeRequest(t, APIRequest{
|
|
Method: "DELETE",
|
|
Path: "/api/governance/virtual-keys/" + vkID,
|
|
})
|
|
|
|
if deleteResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to delete VK: status %d", deleteResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("VK deleted")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify VK and all related resources are removed from memory
|
|
getDataResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap2 := getDataResp2.Body["virtual_keys"].(map[string]interface{})
|
|
|
|
getBudgetsResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
budgetsMap2 := getBudgetsResp2.Body["budgets"].(map[string]interface{})
|
|
|
|
getRateLimitsResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/rate-limits?from_memory=true",
|
|
})
|
|
|
|
rateLimitsMap2 := getRateLimitsResp2.Body["rate_limits"].(map[string]interface{})
|
|
|
|
// VK should be gone
|
|
_, vkStillExists := virtualKeysMap2[vkValue]
|
|
if vkStillExists {
|
|
t.Fatalf("VK still exists in memory after deletion")
|
|
}
|
|
|
|
// Budgets should be gone
|
|
_, vkBudgetStillExists := budgetsMap2[vkBudgetID]
|
|
_, pcBudgetStillExists := budgetsMap2[pcBudgetID]
|
|
if vkBudgetStillExists || pcBudgetStillExists {
|
|
t.Fatalf("Budgets should be cascade-deleted: VK budget exists=%v, PC budget exists=%v",
|
|
vkBudgetStillExists, pcBudgetStillExists)
|
|
}
|
|
|
|
// Rate limits should be gone
|
|
_, vkRateLimitStillExists := rateLimitsMap2[vkRateLimitID]
|
|
_, pcRateLimitStillExists := rateLimitsMap2[pcRateLimitID]
|
|
if vkRateLimitStillExists || pcRateLimitStillExists {
|
|
t.Logf("Note: Rate limits may still exist in memory (orphaned) - this is acceptable")
|
|
}
|
|
|
|
t.Logf("VK removed from memory after deletion ✓")
|
|
t.Logf("VK deletion cascade verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 9: Team/Customer Deletion Should Delete Budget
|
|
// ============================================================================
|
|
|
|
// TestTeamDeletionDeletesBudget verifies that deleting a team also deletes its budget from memory
|
|
func TestTeamDeletionDeletesBudget(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create team with budget
|
|
teamName := "test-team-delete-budget-" + generateRandomID()
|
|
createTeamResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: teamName,
|
|
Budgets: []BudgetRequest{{
|
|
MaxLimit: 100.0,
|
|
ResetDuration: "1h",
|
|
}},
|
|
},
|
|
})
|
|
|
|
if createTeamResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team: status %d", createTeamResp.StatusCode)
|
|
}
|
|
|
|
teamID := ExtractIDFromResponse(t, createTeamResp)
|
|
// Don't add to testData since we'll delete manually
|
|
|
|
t.Logf("Created team with budget")
|
|
|
|
// Get budget ID from in-memory store
|
|
getTeamsResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/teams?from_memory=true",
|
|
})
|
|
|
|
teamsMap1 := getTeamsResp1.Body["teams"].(map[string]interface{})
|
|
|
|
getBudgetsResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
budgetsMap1 := getBudgetsResp1.Body["budgets"].(map[string]interface{})
|
|
|
|
teamData1 := teamsMap1[teamID].(map[string]interface{})
|
|
// Teams now expose a `budgets` array instead of a single `budget_id`.
|
|
budgetsList, ok := teamData1["budgets"].([]interface{})
|
|
if !ok || len(budgetsList) == 0 {
|
|
t.Fatalf("Team %s has no budgets in memory before deletion", teamID)
|
|
}
|
|
budgetID := budgetsList[0].(map[string]interface{})["id"].(string)
|
|
|
|
_, budgetExists := budgetsMap1[budgetID]
|
|
if !budgetExists {
|
|
t.Fatalf("Budget not found in memory before deletion")
|
|
}
|
|
|
|
t.Logf("Team and budget exist in memory ✓")
|
|
|
|
// Delete team
|
|
deleteResp := MakeRequest(t, APIRequest{
|
|
Method: "DELETE",
|
|
Path: "/api/governance/teams/" + teamID,
|
|
})
|
|
|
|
if deleteResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to delete team: status %d", deleteResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Team deleted")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify team and budget are removed from memory
|
|
getTeamsResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/teams?from_memory=true",
|
|
})
|
|
|
|
teamsMap2 := getTeamsResp2.Body["teams"].(map[string]interface{})
|
|
|
|
_, teamStillExists := teamsMap2[teamID]
|
|
if teamStillExists {
|
|
t.Fatalf("Team still exists in memory after deletion")
|
|
}
|
|
|
|
t.Logf("Team removed from memory ✓")
|
|
|
|
// Verify budget is also removed from memory
|
|
getBudgetsResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
if getBudgetsResp2.StatusCode != 200 {
|
|
t.Fatalf("Failed to get budgets from memory: status %d", getBudgetsResp2.StatusCode)
|
|
}
|
|
|
|
budgetsMap2 := getBudgetsResp2.Body["budgets"].(map[string]interface{})
|
|
|
|
_, budgetStillExists := budgetsMap2[budgetID]
|
|
if budgetStillExists {
|
|
t.Fatalf("Budget %s still exists in memory after team deletion", budgetID)
|
|
}
|
|
|
|
t.Logf("Budget removed from memory ✓")
|
|
t.Logf("Team deletion with budget verified ✓")
|
|
}
|
|
|
|
// TestCustomerDeletionDeletesBudget verifies that deleting a customer also deletes its budget from memory
|
|
func TestCustomerDeletionDeletesBudget(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create customer with budget
|
|
customerName := "test-customer-delete-budget-" + generateRandomID()
|
|
createCustomerResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customerName,
|
|
Budget: &BudgetRequest{
|
|
MaxLimit: 100.0,
|
|
ResetDuration: "1h",
|
|
},
|
|
},
|
|
})
|
|
|
|
if createCustomerResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer: status %d", createCustomerResp.StatusCode)
|
|
}
|
|
|
|
customerID := ExtractIDFromResponse(t, createCustomerResp)
|
|
// Don't add to testData since we'll delete manually
|
|
|
|
t.Logf("Created customer with budget")
|
|
|
|
// Get budget ID from in-memory store
|
|
getCustomersResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/customers?from_memory=true",
|
|
})
|
|
|
|
customersMap1 := getCustomersResp1.Body["customers"].(map[string]interface{})
|
|
|
|
getBudgetsResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
budgetsMap1 := getBudgetsResp1.Body["budgets"].(map[string]interface{})
|
|
|
|
customerData1 := customersMap1[customerID].(map[string]interface{})
|
|
budgetID := customerData1["budget_id"].(string)
|
|
|
|
_, budgetExists := budgetsMap1[budgetID]
|
|
if !budgetExists {
|
|
t.Fatalf("Budget not found in memory before deletion")
|
|
}
|
|
|
|
t.Logf("Customer and budget exist in memory ✓")
|
|
|
|
// Delete customer
|
|
deleteResp := MakeRequest(t, APIRequest{
|
|
Method: "DELETE",
|
|
Path: "/api/governance/customers/" + customerID,
|
|
})
|
|
|
|
if deleteResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to delete customer: status %d", deleteResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Customer deleted")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify customer is removed from memory
|
|
getCustomersResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/customers?from_memory=true",
|
|
})
|
|
|
|
customersMap2 := getCustomersResp2.Body["customers"].(map[string]interface{})
|
|
|
|
_, customerStillExists := customersMap2[customerID]
|
|
if customerStillExists {
|
|
t.Fatalf("Customer still exists in memory after deletion")
|
|
}
|
|
|
|
t.Logf("Customer removed from memory ✓")
|
|
|
|
// Verify budget is also removed from memory
|
|
getBudgetsResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/budgets?from_memory=true",
|
|
})
|
|
|
|
if getBudgetsResp2.StatusCode != 200 {
|
|
t.Fatalf("Failed to get budgets from memory: status %d", getBudgetsResp2.StatusCode)
|
|
}
|
|
|
|
budgetsMap2 := getBudgetsResp2.Body["budgets"].(map[string]interface{})
|
|
|
|
_, budgetStillExists := budgetsMap2[budgetID]
|
|
if budgetStillExists {
|
|
t.Fatalf("Budget still exists in memory after customer deletion")
|
|
}
|
|
|
|
t.Logf("Budget removed from memory ✓")
|
|
t.Logf("Customer deletion with budget verified ✓")
|
|
}
|
|
|
|
// ============================================================================
|
|
// SCENARIO 10: Team/Customer Deletion Sets VK entity_id = nil
|
|
// ============================================================================
|
|
|
|
// TestTeamDeletionSetsVKTeamIDToNil verifies that deleting a team sets team_id=nil on associated VKs
|
|
func TestTeamDeletionSetsVKTeamIDToNil(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create team
|
|
teamName := "test-team-vk-nil-" + generateRandomID()
|
|
createTeamResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/teams",
|
|
Body: CreateTeamRequest{
|
|
Name: teamName,
|
|
},
|
|
})
|
|
|
|
if createTeamResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create team: status %d", createTeamResp.StatusCode)
|
|
}
|
|
|
|
teamID := ExtractIDFromResponse(t, createTeamResp)
|
|
// Don't add to testData since we'll delete manually
|
|
|
|
// Create VK assigned to team
|
|
vkName := "test-vk-team-nil-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
TeamID: &teamID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created team and VK assigned to it")
|
|
|
|
// Verify VK has team_id set
|
|
getDataResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap1 := getDataResp1.Body["virtual_keys"].(map[string]interface{})
|
|
vkData1 := virtualKeysMap1[vkValue].(map[string]interface{})
|
|
|
|
teamIDFromVK1, hasTeamID := vkData1["team_id"].(string)
|
|
if !hasTeamID || teamIDFromVK1 != teamID {
|
|
t.Fatalf("VK team_id not set correctly before team deletion")
|
|
}
|
|
|
|
t.Logf("VK has team_id=%s ✓", teamID)
|
|
|
|
// Delete team
|
|
deleteResp := MakeRequest(t, APIRequest{
|
|
Method: "DELETE",
|
|
Path: "/api/governance/teams/" + teamID,
|
|
})
|
|
|
|
if deleteResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to delete team: status %d", deleteResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Team deleted")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify VK still exists but team_id is nil
|
|
getDataResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap2 := getDataResp2.Body["virtual_keys"].(map[string]interface{})
|
|
|
|
vkData2, vkStillExists := virtualKeysMap2[vkValue].(map[string]interface{})
|
|
if !vkStillExists {
|
|
t.Fatalf("VK should still exist after team deletion")
|
|
}
|
|
|
|
teamIDFromVK2, hasTeamID2 := vkData2["team_id"].(string)
|
|
if hasTeamID2 && teamIDFromVK2 != "" {
|
|
t.Fatalf("VK team_id should be nil after team deletion, got: %s", teamIDFromVK2)
|
|
}
|
|
|
|
t.Logf("VK team_id is now nil ✓")
|
|
t.Logf("Team deletion sets VK team_id to nil verified ✓")
|
|
}
|
|
|
|
// TestCustomerDeletionSetsVKCustomerIDToNil verifies that deleting a customer sets customer_id=nil on associated VKs
|
|
func TestCustomerDeletionSetsVKCustomerIDToNil(t *testing.T) {
|
|
t.Parallel()
|
|
testData := NewGlobalTestData()
|
|
defer testData.Cleanup(t)
|
|
|
|
// Create customer
|
|
customerName := "test-customer-vk-nil-" + generateRandomID()
|
|
createCustomerResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/customers",
|
|
Body: CreateCustomerRequest{
|
|
Name: customerName,
|
|
},
|
|
})
|
|
|
|
if createCustomerResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create customer: status %d", createCustomerResp.StatusCode)
|
|
}
|
|
|
|
customerID := ExtractIDFromResponse(t, createCustomerResp)
|
|
// Don't add to testData since we'll delete manually
|
|
|
|
// Create VK assigned directly to customer
|
|
vkName := "test-vk-customer-nil-" + generateRandomID()
|
|
createVKResp := MakeRequest(t, APIRequest{
|
|
Method: "POST",
|
|
Path: "/api/governance/virtual-keys",
|
|
Body: CreateVirtualKeyRequest{
|
|
Name: vkName,
|
|
CustomerID: &customerID,
|
|
},
|
|
})
|
|
|
|
if createVKResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to create VK: status %d", createVKResp.StatusCode)
|
|
}
|
|
|
|
vkID := ExtractIDFromResponse(t, createVKResp)
|
|
testData.AddVirtualKey(vkID)
|
|
|
|
vk := createVKResp.Body["virtual_key"].(map[string]interface{})
|
|
vkValue := vk["value"].(string)
|
|
|
|
t.Logf("Created customer and VK assigned to it")
|
|
|
|
// Verify VK has customer_id set
|
|
getDataResp1 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap1 := getDataResp1.Body["virtual_keys"].(map[string]interface{})
|
|
vkData1 := virtualKeysMap1[vkValue].(map[string]interface{})
|
|
|
|
customerIDFromVK1, hasCustomerID := vkData1["customer_id"].(string)
|
|
if !hasCustomerID || customerIDFromVK1 != customerID {
|
|
t.Fatalf("VK customer_id not set correctly before customer deletion")
|
|
}
|
|
|
|
t.Logf("VK has customer_id=%s ✓", customerID)
|
|
|
|
// Delete customer
|
|
deleteResp := MakeRequest(t, APIRequest{
|
|
Method: "DELETE",
|
|
Path: "/api/governance/customers/" + customerID,
|
|
})
|
|
|
|
if deleteResp.StatusCode != 200 {
|
|
t.Fatalf("Failed to delete customer: status %d", deleteResp.StatusCode)
|
|
}
|
|
|
|
t.Logf("Customer deleted")
|
|
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify VK still exists but customer_id is nil
|
|
getDataResp2 := MakeRequest(t, APIRequest{
|
|
Method: "GET",
|
|
Path: "/api/governance/virtual-keys?from_memory=true",
|
|
})
|
|
|
|
virtualKeysMap2 := getDataResp2.Body["virtual_keys"].(map[string]interface{})
|
|
|
|
vkData2, vkStillExists := virtualKeysMap2[vkValue].(map[string]interface{})
|
|
if !vkStillExists {
|
|
t.Fatalf("VK should still exist after customer deletion")
|
|
}
|
|
|
|
customerIDFromVK2, hasCustomerID2 := vkData2["customer_id"].(string)
|
|
if hasCustomerID2 && customerIDFromVK2 != "" {
|
|
t.Fatalf("VK customer_id should be nil after customer deletion, got: %s", customerIDFromVK2)
|
|
}
|
|
|
|
t.Logf("VK customer_id is now nil ✓")
|
|
t.Logf("Customer deletion sets VK customer_id to nil verified ✓")
|
|
}
|