ragflow/internal/agent/tool/wencai_test.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package tool

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"
)

func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name string
		args string
	}{
		{name: "query", args: `{"query":"商业航天"}`},
		{name: "empty query", args: `{"query":""}`},
		{name: "missing query", args: `{}`},
		{name: "empty arguments", args: ``},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			out, err := NewWencaiTool().InvokableRun(context.Background(), tc.args)
			if err != nil {
				t.Fatalf("InvokableRun errored: %v (out=%s)", err, out)
			}
			var env wencaiEnvelope
			if err := json.Unmarshal([]byte(out), &env); err != nil {
				t.Fatalf("output is not valid JSON: %v (raw=%s)", err, out)
			}
			if env.Report != "" {
				t.Fatalf("report = %q, want empty string", env.Report)
			}
			if env.Error != "" {
				t.Fatalf("_ERROR = %q, want empty", env.Error)
			}
		})
	}
}

func TestWencai_RejectsMalformedJSON(t *testing.T) {
	t.Parallel()

	out, err := NewWencaiTool().InvokableRun(context.Background(), `{not json`)
	if err == nil {
		t.Fatal("expected malformed JSON error")
	}
	if !strings.Contains(err.Error(), "parse arguments") {
		t.Fatalf("err = %q, want parse arguments", err.Error())
	}
	var env wencaiEnvelope
	if jsonErr := json.Unmarshal([]byte(out), &env); jsonErr != nil {
		t.Fatalf("output is not valid JSON: %v (raw=%s)", jsonErr, out)
	}
	if env.Error == "" {
		t.Fatal("_ERROR is empty, want parse error")
	}
}

func TestWencai_RespectsCanceledContext(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	out, err := NewWencaiTool().InvokableRun(ctx, `{"query":"商业航天"}`)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("err = %v, want context.Canceled", err)
	}
	var env wencaiEnvelope
	if jsonErr := json.Unmarshal([]byte(out), &env); jsonErr != nil {
		t.Fatalf("output is not valid JSON: %v (raw=%s)", jsonErr, out)
	}
	if env.Error != context.Canceled.Error() {
		t.Fatalf("_ERROR = %q, want %q", env.Error, context.Canceled.Error())
	}
}

func TestWencai_InfoMatchesPythonMeta(t *testing.T) {
	t.Parallel()

	info, err := NewWencaiTool().Info(context.Background())
	if err != nil {
		t.Fatalf("Info: %v", err)
	}
	if info.Name != "iwencai" {
		t.Fatalf("Name = %q, want iwencai", info.Name)
	}
	if strings.Contains(info.Desc, "STUB") || strings.Contains(info.Desc, "unsupported") {
		t.Fatalf("Desc still exposes stub behavior: %q", info.Desc)
	}
	schema, err := info.ParamsOneOf.ToJSONSchema()
	if err != nil {
		t.Fatalf("ToJSONSchema: %v", err)
	}
	raw, err := json.Marshal(schema)
	if err != nil {
		t.Fatalf("marshal schema: %v", err)
	}
	params := string(raw)
	if !strings.Contains(params, `"query"`) || !strings.Contains(params, `"required":["query"]`) {
		t.Fatalf("schema does not require query: %s", params)
	}
	if strings.Contains(params, `"top_n"`) || strings.Contains(params, `"query_type"`) {
		t.Fatalf("schema leaked node parameters: %s", params)
	}
}

func TestWencai_BuildByNameAcceptsNodeParams(t *testing.T) {
	t.Parallel()

	built, err := BuildByName("wencai", map[string]any{
		"top_n":      float64(20),
		"query_type": "fund",
	})
	if err != nil {
		t.Fatalf("BuildByName: %v", err)
	}
	wencai, ok := built.(*WencaiTool)
	if !ok {
		t.Fatalf("built type = %T, want *WencaiTool", built)
	}
	if wencai.defaults.TopN != 20 {
		t.Fatalf("defaults.TopN = %d, want 20", wencai.defaults.TopN)
	}
	if wencai.defaults.QueryType != "fund" {
		t.Fatalf("defaults.QueryType = %q, want fund", wencai.defaults.QueryType)
	}
}

func TestWencai_BuildByNameUsesPythonDefaults(t *testing.T) {
	t.Parallel()

	built, err := BuildByName("wencai", nil)
	if err != nil {
		t.Fatalf("BuildByName: %v", err)
	}
	wencai := built.(*WencaiTool)
	if wencai.defaults.TopN != 10 || wencai.defaults.QueryType != "stock" {
		t.Fatalf("defaults = %+v, want top_n=10 query_type=stock", wencai.defaults)
	}
}

func TestWencai_BuildByNameRejectsInvalidNodeParams(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name   string
		params map[string]any
	}{
		{name: "zero top_n", params: map[string]any{"top_n": 0}},
		{name: "fractional top_n", params: map[string]any{"top_n": 1.5}},
		{name: "string top_n", params: map[string]any{"top_n": "10"}},
		{name: "invalid query_type", params: map[string]any{"query_type": "crypto"}},
		{name: "non-string query_type", params: map[string]any{"query_type": 1}},
		{name: "unknown key", params: map[string]any{"page": 1}},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if _, err := BuildByName("wencai", tc.params); err == nil {
				t.Fatalf("BuildByName(%v) succeeded, want validation error", tc.params)
			}
		})
	}
}

func TestWencai_MergeDefaults(t *testing.T) {
	t.Parallel()

	got := mergeWencaiParams(
		wencaiParams{Query: "configured", TopN: 20, QueryType: "fund"},
		wencaiParams{Query: "商业航天"},
	)
	if got.Query != "商业航天" || got.TopN != 20 || got.QueryType != "fund" {
		t.Fatalf("merged params = %+v", got)
	}
}