yao/agent/search/nlp/querydsl/generator_test.go

package querydsl

import (
	"testing"

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

func TestNewGenerator(t *testing.T) {
	tests := []struct {
		name         string
		usesQueryDSL string
		config       *types.QueryDSLConfig
	}{
		{
			name:         "builtin mode",
			usesQueryDSL: "builtin",
			config:       nil,
		},
		{
			name:         "empty defaults to builtin",
			usesQueryDSL: "",
			config:       nil,
		},
		{
			name:         "agent mode",
			usesQueryDSL: "my-querydsl-agent",
			config:       &types.QueryDSLConfig{Strict: true},
		},
		{
			name:         "mcp mode",
			usesQueryDSL: "mcp:nlp.generate_querydsl",
			config:       nil,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gen := NewGenerator(tt.usesQueryDSL, tt.config)
			assert.NotNil(t, gen)
			assert.Equal(t, tt.usesQueryDSL, gen.usesQueryDSL)
			assert.Equal(t, tt.config, gen.config)
		})
	}
}

func TestGenerator_Generate_Builtin_RequiresContext(t *testing.T) {
	// Builtin mode now uses __yao.querydsl agent which requires context
	gen := NewGenerator("builtin", nil)

	input := &Input{
		Query:    "find all active users",
		ModelIDs: []string{"user"},
		Limit:    10,
	}

	// Without context, should return error
	_, err := gen.Generate(nil, input)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "context is required")
}

func TestGenerator_Generate_EmptyMode_RequiresContext(t *testing.T) {
	// Empty mode defaults to __yao.querydsl agent which requires context
	gen := NewGenerator("", nil)

	input := &Input{
		Query:    "search products",
		ModelIDs: []string{"product"},
		Limit:    5,
	}

	// Without context, should return error
	_, err := gen.Generate(nil, input)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "context is required")
}

func TestGenerator_Generate_AgentMode_RequiresContext(t *testing.T) {
	// Custom agent mode requires context
	gen := NewGenerator("custom.querydsl.agent", nil)

	input := &Input{
		Query:    "find users",
		ModelIDs: []string{"user"},
		Limit:    10,
	}

	_, err := gen.Generate(nil, input)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "context is required")
}

func TestGenerator_Generate_MCPMode_InvalidFormat(t *testing.T) {
	// Invalid MCP format should fallback to system agent (which requires context)
	gen := NewGenerator("mcp:invalid", nil)

	input := &Input{
		Query:    "find users",
		ModelIDs: []string{"user"},
		Limit:    10,
	}

	_, err := gen.Generate(nil, input)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "context is required")
}

func TestSystemQueryDSLAgentConstant(t *testing.T) {
	// Verify the system querydsl agent constant
	assert.Equal(t, "__yao.querydsl", SystemQueryDSLAgent)
}

func TestResult(t *testing.T) {
	result := &Result{
		DSL: &gou.QueryDSL{
			Limit: 10,
		},
		Explain:  "Generated query for finding users",
		Warnings: []string{"using placeholder implementation"},
	}

	assert.NotNil(t, result.DSL)
	assert.Equal(t, 10, result.DSL.Limit)
	assert.NotEmpty(t, result.Explain)
	assert.Len(t, result.Warnings, 1)
}

func TestGenerator_ValidateFields(t *testing.T) {
	gen := NewGenerator("", nil)

	t.Run("validate select fields", func(t *testing.T) {
		result := &Result{
			DSL: &gou.QueryDSL{
				Select: []gou.Expression{
					{Field: "id"},
					{Field: "name"},
					{Field: "secret_field"},
				},
			},
		}
		allowedFields := []string{"id", "name", "email"}

		validated := gen.validateFields(result, allowedFields)
		assert.NotNil(t, validated)
		assert.Len(t, validated.DSL.Select, 2)
		assert.Contains(t, validated.Warnings[0], "secret_field")
	})

	t.Run("validate where fields", func(t *testing.T) {
		result := &Result{
			DSL: &gou.QueryDSL{
				Wheres: []gou.Where{
					{
						Condition: gou.Condition{
							Field: &gou.Expression{Field: "status"},
							OP:    "=",
							Value: "active",
						},
					},
					{
						Condition: gou.Condition{
							Field: &gou.Expression{Field: "secret"},
							OP:    "=",
							Value: "hidden",
						},
					},
				},
			},
		}
		allowedFields := []string{"status", "name"}

		validated := gen.validateFields(result, allowedFields)
		assert.NotNil(t, validated)
		assert.Len(t, validated.DSL.Wheres, 1)
		assert.Contains(t, validated.Warnings[0], "secret")
	})

	t.Run("validate order fields", func(t *testing.T) {
		result := &Result{
			DSL: &gou.QueryDSL{
				Orders: gou.Orders{
					{Field: &gou.Expression{Field: "created_at"}, Sort: "desc"},
					{Field: &gou.Expression{Field: "secret_sort"}, Sort: "asc"},
				},
			},
		}
		allowedFields := []string{"created_at", "updated_at"}

		validated := gen.validateFields(result, allowedFields)
		assert.NotNil(t, validated)
		assert.Len(t, validated.DSL.Orders, 1)
		assert.Contains(t, validated.Warnings[0], "secret_sort")
	})

	t.Run("nil DSL", func(t *testing.T) {
		result := &Result{DSL: nil}
		allowedFields := []string{"id", "name"}

		validated := gen.validateFields(result, allowedFields)
		assert.NotNil(t, validated)
		assert.Nil(t, validated.DSL)
	})
}