lijiaoqiao/gateway/internal/router/engine/routing_engine_test.go

package engine

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"

	"lijiaoqiao/gateway/internal/adapter"
	"lijiaoqiao/gateway/internal/router/strategy"
)

// TestRoutingEngine_SelectProvider 测试路由引擎根据策略选择provider
func TestRoutingEngine_SelectProvider(t *testing.T) {
	engine := NewRoutingEngine()

	// 注册策略
	costBased := strategy.NewCostBasedTemplate("CostBased", strategy.CostParams{
		MaxCostPer1KTokens: 1.0,
	})

	// 注册providers
	costBased.RegisterProvider("ProviderA", &MockProvider{
		name:            "ProviderA",
		costPer1KTokens: 0.5,
		available:       true,
		models:          []string{"gpt-4"},
	})
	costBased.RegisterProvider("ProviderB", &MockProvider{
		name:            "ProviderB",
		costPer1KTokens: 0.3, // 最低成本
		available:       true,
		models:          []string{"gpt-4"},
	})

	engine.RegisterStrategy("cost_based", costBased)

	req := &strategy.RoutingRequest{
		Model:   "gpt-4",
		UserID:  "user123",
		MaxCost: 1.0,
	}

	decision, err := engine.SelectProvider(context.Background(), req, "cost_based")

	assert.NoError(t, err)
	assert.NotNil(t, decision)
	assert.Equal(t, "ProviderB", decision.Provider, "Should select lowest cost provider")
	assert.True(t, decision.TakeoverMark, "TakeoverMark should be true for M-008")
}

// TestRoutingEngine_DecisionMetrics 测试路由决策记录metrics
func TestRoutingEngine_DecisionMetrics(t *testing.T) {
	engine := NewRoutingEngine()

	// 创建mock metrics collector
	engine.metrics = &MockRoutingMetrics{}

	// 注册策略
	costBased := strategy.NewCostBasedTemplate("CostBased", strategy.CostParams{
		MaxCostPer1KTokens: 1.0,
	})

	costBased.RegisterProvider("ProviderA", &MockProvider{
		name:            "ProviderA",
		costPer1KTokens: 0.5,
		available:       true,
		models:          []string{"gpt-4"},
	})

	engine.RegisterStrategy("cost_based", costBased)

	req := &strategy.RoutingRequest{
		Model:  "gpt-4",
		UserID: "user123",
	}

	decision, err := engine.SelectProvider(context.Background(), req, "cost_based")

	assert.NoError(t, err)
	assert.NotNil(t, decision)

	// 验证metrics被记录
	metrics := engine.metrics.(*MockRoutingMetrics)
	assert.True(t, metrics.recordCalled, "RecordSelection should be called")
	assert.Equal(t, "ProviderA", metrics.lastProvider, "Provider should be recorded")
}

// MockProvider 用于测试的Mock Provider
type MockProvider struct {
	name            string
	costPer1KTokens float64
	qualityScore    float64
	latencyMs       int64
	available       bool
	models          []string
}

func (m *MockProvider) ChatCompletion(ctx context.Context, model string, messages []adapter.Message, options adapter.CompletionOptions) (*adapter.CompletionResponse, error) {
	return nil, nil
}

func (m *MockProvider) ChatCompletionStream(ctx context.Context, model string, messages []adapter.Message, options adapter.CompletionOptions) (<-chan *adapter.StreamChunk, error) {
	return nil, nil
}

func (m *MockProvider) GetUsage(response *adapter.CompletionResponse) adapter.Usage {
	return adapter.Usage{}
}

func (m *MockProvider) MapError(err error) adapter.ProviderError {
	return adapter.ProviderError{}
}

func (m *MockProvider) HealthCheck(ctx context.Context) bool {
	return m.available
}

func (m *MockProvider) ProviderName() string {
	return m.name
}

func (m *MockProvider) SupportedModels() []string {
	return m.models
}

func (m *MockProvider) GetCostPer1KTokens() float64 {
	return m.costPer1KTokens
}

func (m *MockProvider) GetQualityScore() float64 {
	return m.qualityScore
}

func (m *MockProvider) GetLatencyMs() int64 {
	return m.latencyMs
}

// MockRoutingMetrics 用于测试的Mock Metrics
type MockRoutingMetrics struct {
	recordCalled  bool
	lastProvider  string
	lastStrategy  string
	takeoverMark  bool
}

func (m *MockRoutingMetrics) RecordSelection(provider string, strategyName string, decision *strategy.RoutingDecision) {
	m.recordCalled = true
	m.lastProvider = provider
	m.lastStrategy = strategyName
	if decision != nil {
		m.takeoverMark = decision.TakeoverMark
	}
}