yao/agent/search/handlers/kb/handler_test.go

package kb

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/yaoapp/yao/agent/search/types"
)

func TestNewHandler(t *testing.T) {
	t.Run("with nil config", func(t *testing.T) {
		h := NewHandler(nil)
		assert.NotNil(t, h)
		assert.Nil(t, h.config)
	})

	t.Run("with config", func(t *testing.T) {
		cfg := &types.KBConfig{
			Collections: []string{"docs", "faq"},
			Threshold:   0.8,
			Graph:       true,
		}
		h := NewHandler(cfg)
		assert.NotNil(t, h)
		assert.Equal(t, cfg, h.config)
	})
}

func TestHandler_Type(t *testing.T) {
	h := NewHandler(nil)
	assert.Equal(t, types.SearchTypeKB, h.Type())
}

func TestHandler_Search_Validation(t *testing.T) {
	tests := []struct {
		name        string
		config      *types.KBConfig
		req         *types.Request
		expectError string
	}{
		{
			name:   "empty query",
			config: nil,
			req: &types.Request{
				Type:  types.SearchTypeKB,
				Query: "",
			},
			expectError: "query is required",
		},
		{
			name:   "no collections - KB not initialized",
			config: nil,
			req: &types.Request{
				Type:  types.SearchTypeKB,
				Query: "test query",
			},
			expectError: "knowledge base not initialized",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			h := NewHandler(tt.config)
			result, err := h.Search(tt.req)

			assert.NoError(t, err)
			assert.NotNil(t, result)
			assert.Equal(t, types.SearchTypeKB, result.Type)
			assert.Equal(t, tt.req.Query, result.Query)

			if tt.expectError != "" {
				assert.Equal(t, tt.expectError, result.Error)
			} else {
				assert.Empty(t, result.Error)
			}

			// Duration should be set
			assert.GreaterOrEqual(t, result.Duration, int64(0))
		})
	}
}

func TestHandler_Search_SourcePreserved(t *testing.T) {
	h := NewHandler(&types.KBConfig{Collections: []string{"docs"}})

	sources := []types.SourceType{types.SourceUser, types.SourceHook, types.SourceAuto}
	for _, source := range sources {
		req := &types.Request{
			Type:        types.SearchTypeKB,
			Query:       "test",
			Source:      source,
			Collections: []string{"docs"},
		}
		result, err := h.Search(req)
		assert.NoError(t, err)
		assert.Equal(t, source, result.Source)
	}
}

func TestHandler_Search_CollectionsFromConfig(t *testing.T) {
	cfg := &types.KBConfig{
		Collections: []string{"docs", "faq"},
		Threshold:   0.7,
	}
	h := NewHandler(cfg)

	// Request without collections should use config collections
	req := &types.Request{
		Type:  types.SearchTypeKB,
		Query: "test query",
	}
	result, err := h.Search(req)

	assert.NoError(t, err)
	assert.NotNil(t, result)
	// Without KB initialized, we get "knowledge base not initialized" error
	assert.Equal(t, "knowledge base not initialized", result.Error)
}

func TestHandler_Search_CollectionsFromRequest(t *testing.T) {
	h := NewHandler(nil)

	// Request with collections
	req := &types.Request{
		Type:        types.SearchTypeKB,
		Query:       "test query",
		Collections: []string{"docs", "faq"},
	}
	result, err := h.Search(req)

	assert.NoError(t, err)
	assert.NotNil(t, result)
	// Without KB initialized, we get "knowledge base not initialized" error
	assert.Equal(t, "knowledge base not initialized", result.Error)
}

func TestHandler_Search_ThresholdHandling(t *testing.T) {
	tests := []struct {
		name            string
		configThreshold float64
		reqThreshold    float64
	}{
		{
			name:            "threshold from request",
			configThreshold: 0.7,
			reqThreshold:    0.9,
		},
		{
			name:            "threshold from config",
			configThreshold: 0.8,
			reqThreshold:    0,
		},
		{
			name:            "default threshold",
			configThreshold: 0,
			reqThreshold:    0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var cfg *types.KBConfig
			if tt.configThreshold > 0 {
				cfg = &types.KBConfig{
					Collections: []string{"docs"},
					Threshold:   tt.configThreshold,
				}
			}
			h := NewHandler(cfg)

			req := &types.Request{
				Type:        types.SearchTypeKB,
				Query:       "test query",
				Threshold:   tt.reqThreshold,
				Collections: []string{"docs"},
			}
			result, err := h.Search(req)

			assert.NoError(t, err)
			assert.NotNil(t, result)
		})
	}
}

func TestHandler_Search_GraphMode(t *testing.T) {
	tests := []struct {
		name        string
		configGraph bool
		reqGraph    bool
	}{
		{
			name:        "graph from request",
			configGraph: false,
			reqGraph:    true,
		},
		{
			name:        "graph from config",
			configGraph: true,
			reqGraph:    false,
		},
		{
			name:        "no graph",
			configGraph: false,
			reqGraph:    false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cfg := &types.KBConfig{
				Collections: []string{"docs"},
				Graph:       tt.configGraph,
			}
			h := NewHandler(cfg)

			req := &types.Request{
				Type:        types.SearchTypeKB,
				Query:       "test query",
				Collections: []string{"docs"},
				Graph:       tt.reqGraph,
			}
			result, err := h.Search(req)

			assert.NoError(t, err)
			assert.NotNil(t, result)
		})
	}
}

func TestHandler_Search_LimitHandling(t *testing.T) {
	h := NewHandler(&types.KBConfig{Collections: []string{"docs"}})

	tests := []struct {
		name  string
		limit int
	}{
		{
			name:  "custom limit",
			limit: 5,
		},
		{
			name:  "default limit",
			limit: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := &types.Request{
				Type:        types.SearchTypeKB,
				Query:       "test query",
				Collections: []string{"docs"},
				Limit:       tt.limit,
			}
			result, err := h.Search(req)

			assert.NoError(t, err)
			assert.NotNil(t, result)
		})
	}
}