yao/agent/llm/resolve_test.go

package llm_test

import (
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/yaoapp/gou/connector"
	"github.com/yaoapp/gou/store"
	"github.com/yaoapp/yao/agent/llm"
	"github.com/yaoapp/yao/config"
	"github.com/yaoapp/yao/llmprovider"
	"github.com/yaoapp/yao/setting"
	"github.com/yaoapp/yao/test"
)

func TestMain(m *testing.M) {
	test.Prepare(nil, config.Conf)
	defer test.Clean()
	os.Exit(m.Run())
}

type mockIdentity struct {
	UserID string
	TeamID string
}

func (m *mockIdentity) GetUserID() string { return m.UserID }
func (m *mockIdentity) GetTeamID() string { return m.TeamID }

func setupResolveTest(t *testing.T) string {
	t.Helper()
	test.Prepare(t, config.Conf)

	err := setting.Init()
	require.NoError(t, err)

	err = llmprovider.Init()
	require.NoError(t, err)

	connIDs := connector.AIConnectors
	if len(connIDs) == 0 {
		t.Skip("no AI connectors available in test env")
	}

	cid := connIDs[0].Value

	t.Cleanup(func() {
		s, _ := store.Get("__yao.store")
		if s != nil {
			s.Del("llmprovider:*")
		}
		c, _ := store.Get("__yao.cache")
		if c != nil {
			c.Del("llmprovider:*")
		}
		test.Clean()
	})

	return cid
}

// --- use:: prefix tests ---

func TestResolveConnector_UseLight(t *testing.T) {
	cid := setupResolveTest(t)

	err := llmprovider.Global.SetDefaults(map[string]string{
		"default": cid,
		"light":   cid,
	})
	require.NoError(t, err)

	conn, caps, err := llm.ResolveConnector("use::light", nil)
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}

func TestResolveConnector_UseDefault(t *testing.T) {
	cid := setupResolveTest(t)

	err := llmprovider.Global.SetDefaults(map[string]string{
		"default": cid,
	})
	require.NoError(t, err)

	conn, caps, err := llm.ResolveConnector("use::default", nil)
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}

func TestResolveConnector_UseLightWithIdentity(t *testing.T) {
	cid := setupResolveTest(t)

	err := llmprovider.Global.SetDefaults(map[string]string{
		"default": cid,
		"light":   cid,
	})
	require.NoError(t, err)

	conn, caps, err := llm.ResolveConnector("use::light", &mockIdentity{UserID: "u1", TeamID: "t1"})
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}

func TestResolveConnector_UseLightNoProvider(t *testing.T) {
	test.Prepare(t, config.Conf)
	defer test.Clean()

	saved := llmprovider.Global
	llmprovider.Global = nil
	defer func() { llmprovider.Global = saved }()

	_, _, err := llm.ResolveConnector("use::light", nil)
	assert.Error(t, err)
}

// --- Explicit connector tests ---

func TestResolveConnector_ExplicitID(t *testing.T) {
	cid := setupResolveTest(t)

	conn, caps, err := llm.ResolveConnector(cid, nil)
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}

func TestResolveConnector_ExplicitIDPriority(t *testing.T) {
	cid := setupResolveTest(t)

	err := llmprovider.Global.SetDefaults(map[string]string{
		"default": cid,
		"light":   cid,
	})
	require.NoError(t, err)

	// Explicit connector ID is NOT a use:: prefix, so it takes priority
	conn, caps, err := llm.ResolveConnector(cid, &mockIdentity{UserID: "u1"})
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}

func TestResolveConnector_InvalidID(t *testing.T) {
	setupResolveTest(t)

	_, _, err := llm.ResolveConnector("nonexistent-connector-xyz", nil)
	assert.Error(t, err)
}

// --- Empty connector fallback ---

func TestResolveConnector_EmptyFallbackDefault(t *testing.T) {
	cid := setupResolveTest(t)

	err := llmprovider.Global.SetDefaults(map[string]string{
		"default": cid,
	})
	require.NoError(t, err)

	// Empty string → treated as use::default
	conn, caps, err := llm.ResolveConnector("", nil)
	require.NoError(t, err)
	assert.NotNil(t, conn)
	assert.NotNil(t, caps)
}