bda75aeb0d
- Updated .env.example to include new configuration options for AI gateway and WebUI secret key. - Modified Nginx configuration to support additional API routes for model management and migration. - Implemented new API endpoints for discovering organization-level LLM models and managing hosted mail services. - Enhanced AI gateway logic to support organization-specific model access and permissions. - Improved error handling and response structures in the AI and mail APIs. - Added integration tests for new features and updated existing tests for model access control.
56 lines
1.5 KiB
Go
56 lines
1.5 KiB
Go
package ai
|
|
|
|
import "testing"
|
|
|
|
func TestApplyModelCatalogEmptyAllowsAll(t *testing.T) {
|
|
upstream := []map[string]any{
|
|
{"id": "gpt-4o", "object": "model"},
|
|
{"id": "llama3", "object": "model"},
|
|
}
|
|
got := ApplyModelCatalog(upstream, nil)
|
|
if len(got) != 2 {
|
|
t.Fatalf("expected 2 models, got %d", len(got))
|
|
}
|
|
}
|
|
|
|
func TestApplyModelCatalogAllowlist(t *testing.T) {
|
|
upstream := []map[string]any{
|
|
{"id": "gpt-4o", "object": "model", "owned_by": "OpenAI"},
|
|
{"id": "llama3", "object": "model", "owned_by": "Local"},
|
|
}
|
|
catalog := []ModelCatalogEntry{
|
|
{ModelID: "gpt-4o", Label: "GPT-4o", Enabled: true},
|
|
{ModelID: "llama3", Label: "Llama 3", Enabled: false},
|
|
{ModelID: "manual-model", Label: "Manuel", Enabled: true},
|
|
}
|
|
got := ApplyModelCatalog(upstream, catalog)
|
|
if len(got) != 2 {
|
|
t.Fatalf("expected 2 allowed models, got %d", len(got))
|
|
}
|
|
if got[0]["label"] != "GPT-4o" {
|
|
t.Fatalf("expected label GPT-4o, got %v", got[0]["label"])
|
|
}
|
|
if got[1]["id"] != "manual-model" {
|
|
t.Fatalf("expected manual-model, got %v", got[1]["id"])
|
|
}
|
|
}
|
|
|
|
func TestIsModelAllowed(t *testing.T) {
|
|
catalog := []ModelCatalogEntry{
|
|
{ModelID: "gpt-4o", Enabled: true},
|
|
{ModelID: "llama3", Enabled: false},
|
|
}
|
|
if !IsModelAllowed("gpt-4o", catalog) {
|
|
t.Fatal("gpt-4o should be allowed")
|
|
}
|
|
if IsModelAllowed("llama3", catalog) {
|
|
t.Fatal("llama3 should be blocked")
|
|
}
|
|
if IsModelAllowed("unknown", catalog) {
|
|
t.Fatal("unknown model should be blocked when catalog is set")
|
|
}
|
|
if !IsModelAllowed("unknown", nil) {
|
|
t.Fatal("unknown model should be allowed without catalog")
|
|
}
|
|
}
|