package llmprovider_test import ( "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/yaoapp/gou/process" "github.com/yaoapp/gou/store" "github.com/yaoapp/yao/llmprovider" "github.com/yaoapp/yao/setting" ) func setupRegistryWithSetting(t *testing.T) *llmprovider.Registry { t.Helper() r := setupRegistry(t) err := setting.Init() require.NoError(t, err) t.Cleanup(func() { s, _ := store.Get("__yao.store") if s != nil { s.Del("setting:*") } c, _ := store.Get("__yao.cache") if c != nil { c.Del("setting:*") } }) return r } func createTestProviderForRole(t *testing.T, r *llmprovider.Registry, key string) *llmprovider.Provider { t.Helper() p := llmprovider.Provider{ Key: key, Name: "Test " + key, Type: "openai", APIURL: "https://api.openai.com", APIKey: "sk-test-role", Enabled: true, Models: []llmprovider.ModelInfo{ {ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls", "streaming"}, Enabled: true}, }, Owner: llmprovider.ProviderOwner{Type: "system"}, } created, err := r.Create(&p) require.NoError(t, err) return created } func TestSetDefaults(t *testing.T) { r := setupRegistryWithSetting(t) p := createTestProviderForRole(t, r, "sd-provider") err := r.SetDefaults(map[string]string{ "default": p.Key, }) require.NoError(t, err) merged, err := setting.Global.GetMerged("", "", "llm.roles") require.NoError(t, err) def, ok := merged["default"].(map[string]interface{}) require.True(t, ok) assert.Equal(t, p.Key, def["provider"]) assert.Equal(t, "gpt-4o", def["model"]) } func TestGetRole(t *testing.T) { r := setupRegistryWithSetting(t) p := createTestProviderForRole(t, r, "role-provider") err := r.SetDefaults(map[string]string{"default": p.Key}) require.NoError(t, err) cid, err := r.GetRole("default") require.NoError(t, err) assert.Equal(t, p.ConnectorID, cid) } func TestGetRoleByUser(t *testing.T) { r := setupRegistryWithSetting(t) sysP := createTestProviderForRole(t, r, "sys-prov") err := r.SetDefaults(map[string]string{"default": sysP.Key}) require.NoError(t, err) userP := createTestProviderForRole(t, r, "user-prov") _, err = setting.Global.Set( setting.ScopeID{Scope: setting.ScopeUser, UserID: "u1"}, "llm.roles", map[string]interface{}{ "default": map[string]interface{}{ "provider": userP.Key, "model": "gpt-4o", }, }, ) require.NoError(t, err) cid, err := r.GetRoleByUser("default", "u1") require.NoError(t, err) assert.Equal(t, userP.ConnectorID, cid, "user scope should override system") cidSys, err := r.GetRole("default") require.NoError(t, err) assert.Equal(t, sysP.ConnectorID, cidSys, "system scope should still return system provider") } func TestGetRoleByTeam(t *testing.T) { r := setupRegistryWithSetting(t) sysP := createTestProviderForRole(t, r, "sys-team-prov") err := r.SetDefaults(map[string]string{"default": sysP.Key}) require.NoError(t, err) teamP := createTestProviderForRole(t, r, "team-prov") _, err = setting.Global.Set( setting.ScopeID{Scope: setting.ScopeTeam, TeamID: "t1"}, "llm.roles", map[string]interface{}{ "default": map[string]interface{}{ "provider": teamP.Key, "model": "gpt-4o", }, }, ) require.NoError(t, err) cid, err := r.GetRoleByTeam("default", "t1") require.NoError(t, err) assert.Equal(t, teamP.ConnectorID, cid, "team scope should override system") } func TestGetRoleNotConfigured(t *testing.T) { _ = setupRegistryWithSetting(t) _, err := llmprovider.Global.GetRole("nonexistent") assert.Error(t, err) assert.Contains(t, err.Error(), "not configured") } func TestListRoles(t *testing.T) { r := setupRegistryWithSetting(t) p := createTestProviderForRole(t, r, "list-prov") err := r.SetDefaults(map[string]string{ "default": p.Key, "vision": p.Key, }) require.NoError(t, err) roles, err := r.ListRoles() require.NoError(t, err) assert.Contains(t, roles, "default") assert.Contains(t, roles, "vision") assert.Equal(t, p.Key, roles["default"].Provider) } func TestListRolesByUser(t *testing.T) { r := setupRegistryWithSetting(t) sysP := createTestProviderForRole(t, r, "list-sys") userP := createTestProviderForRole(t, r, "list-user") err := r.SetDefaults(map[string]string{"default": sysP.Key, "vision": sysP.Key}) require.NoError(t, err) _, err = setting.Global.Set( setting.ScopeID{Scope: setting.ScopeUser, UserID: "u2"}, "llm.roles", map[string]interface{}{ "default": map[string]interface{}{ "provider": userP.Key, "model": "gpt-4o", }, }, ) require.NoError(t, err) roles, err := r.ListRolesByUser("u2") require.NoError(t, err) assert.Equal(t, userP.Key, roles["default"].Provider, "user override for default") assert.Equal(t, sysP.Key, roles["vision"].Provider, "system fallback for vision") } func TestProcessGetRole(t *testing.T) { r := setupRegistryWithSetting(t) p := createTestProviderForRole(t, r, "proc-role") err := r.SetDefaults(map[string]string{"default": p.Key}) require.NoError(t, err) proc := process.New("llmprovider.getrole", "default") result, err := proc.Exec() require.NoError(t, err) assert.Equal(t, p.ConnectorID, result) } func TestProcessListRoles(t *testing.T) { r := setupRegistryWithSetting(t) p := createTestProviderForRole(t, r, "proc-list-role") err := r.SetDefaults(map[string]string{"default": p.Key}) require.NoError(t, err) proc := process.New("llmprovider.listroles") result, err := proc.Exec() require.NoError(t, err) m, ok := result.(map[string]interface{}) require.True(t, ok) def, ok := m["default"].(map[string]interface{}) require.True(t, ok) assert.Equal(t, p.Key, def["provider"]) }