ragflow/internal/agent/tool/tushare_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"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

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

	cases := []struct {
		name         string
		params       tushareParams
		wantToken    string
		wantAPIName  string
		wantFieldKey string
		wantFieldVal string
	}{
		{
			name: "minimal — token + api_name",
			params: tushareParams{
				Token:   "T-abc",
				APIName: "stock_basic",
			},
			wantToken:   "T-abc",
			wantAPIName: "stock_basic",
		},
		{
			name: "with params and fields",
			params: tushareParams{
				Token:   "T-xyz",
				APIName: "daily",
				Params:  map[string]string{"ts_code": "000001.SZ", "start_date": "20240101"},
				Fields:  []string{"ts_code", "trade_date", "close"},
			},
			wantToken:    "T-xyz",
			wantAPIName:  "daily",
			wantFieldKey: "ts_code",
			wantFieldVal: "000001.SZ",
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			body, err := buildTushareRequestBody(tc.params)
			if err != nil {
				t.Fatalf("buildTushareRequestBody: %v", err)
			}
			var got tushareRequest
			if jerr := json.Unmarshal(body, &got); jerr != nil {
				t.Fatalf("request body is not valid JSON: %v (raw=%s)", jerr, body)
			}
			if got.Token != tc.wantToken {
				t.Errorf("token = %q, want %q", got.Token, tc.wantToken)
			}
			if got.APIName != tc.wantAPIName {
				t.Errorf("api_name = %q, want %q", got.APIName, tc.wantAPIName)
			}
			if tc.wantFieldKey != "" {
				if v, ok := got.Params[tc.wantFieldKey]; !ok || v != tc.wantFieldVal {
					t.Errorf("params[%q] = %q (present=%v), want %q", tc.wantFieldKey, v, ok, tc.wantFieldVal)
				}
			}
		})
	}
}

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

	var (
		gotMethod  string
		gotCT      string
		gotBodyRaw []byte
	)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotMethod = r.Method
		gotCT = r.Header.Get("Content-Type")
		gotBodyRaw, _ = io.ReadAll(r.Body)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"code": 0,
			"msg":  "",
			"data": {
				"fields": ["ts_code", "name", "industry"],
				"items":  [
					["000001.SZ", "平安银行", "银行"],
					["600519.SH", "贵州茅台", "白酒"]
				]
			}
		}`))
	}))
	defer srv.Close()

	prev := tushareEndpoint
	tushareEndpoint = srv.URL
	t.Cleanup(func() { tushareEndpoint = prev })

	helper := NewHTTPHelper().WithClient(&http.Client{
		Transport: rewriteHostTransport(srv.URL),
	})
	tool := NewTushareToolWith(helper)

	out, err := tool.InvokableRun(context.Background(),
		`{"token":"T-test","api_name":"stock_basic","params":{"list_status":"L"}}`)
	if err != nil {
		t.Fatalf("InvokableRun: %v", err)
	}
	if gotMethod != http.MethodPost {
		t.Errorf("method = %q, want POST", gotMethod)
	}
	if !strings.Contains(gotCT, "application/json") {
		t.Errorf("Content-Type = %q, want application/json", gotCT)
	}

	// Verify the request body had the right shape.
	var sentReq tushareRequest
	if jerr := json.Unmarshal(gotBodyRaw, &sentReq); jerr != nil {
		t.Fatalf("server saw malformed request body: %v (raw=%s)", jerr, gotBodyRaw)
	}
	if sentReq.Token != "T-test" {
		t.Errorf("server saw token = %q, want T-test", sentReq.Token)
	}
	if sentReq.APIName != "stock_basic" {
		t.Errorf("server saw api_name = %q, want stock_basic", sentReq.APIName)
	}
	if v := sentReq.Params["list_status"]; v != "L" {
		t.Errorf("server saw params[list_status] = %q, want L", v)
	}

	var env tushareEnvelope
	if jerr := json.Unmarshal([]byte(out), &env); jerr != nil {
		t.Fatalf("output is not valid JSON: %v (raw=%s)", jerr, out)
	}
	if env.Error != "" {
		t.Errorf("Error = %q, want empty", env.Error)
	}
	if len(env.Fields) != 3 || env.Fields[0] != "ts_code" {
		t.Errorf("Fields = %v, want [ts_code name industry]", env.Fields)
	}
	if len(env.Items) != 2 {
		t.Fatalf("Items len = %d, want 2", len(env.Items))
	}
	if env.Items[0][1] != "平安银行" {
		t.Errorf("Items[0][1] = %v, want 平安银行", env.Items[0][1])
	}
}

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

	tool := NewTushareTool()
	_, err := tool.InvokableRun(context.Background(),
		`{"api_name":"stock_basic"}`)
	if err == nil {
		t.Fatal("expected error for missing token")
	}
	if !strings.Contains(err.Error(), "token") {
		t.Errorf("err = %v, want to mention token", err)
	}
}

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

	tool := NewTushareTool()
	_, err := tool.InvokableRun(context.Background(),
		`{"token":"T-abc"}`)
	if err == nil {
		t.Fatal("expected error for missing api_name")
	}
	if !strings.Contains(err.Error(), "api_name") {
		t.Errorf("err = %v, want to mention api_name", err)
	}
}

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

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"code": 40201, "msg": "权限不足"}`))
	}))
	defer srv.Close()

	prev := tushareEndpoint
	tushareEndpoint = srv.URL
	t.Cleanup(func() { tushareEndpoint = prev })

	helper := NewHTTPHelper().WithClient(&http.Client{
		Transport: rewriteHostTransport(srv.URL),
	})
	tool := NewTushareToolWith(helper)

	_, err := tool.InvokableRun(context.Background(),
		`{"token":"T-abc","api_name":"premium_only"}`)
	if err == nil {
		t.Fatal("expected error for non-zero code, got nil")
	}
	if !strings.Contains(err.Error(), "权限不足") {
		t.Errorf("err = %v, want to surface upstream msg", err)
	}
}

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

	tool := NewTushareTool()
	info, err := tool.Info(context.Background())
	if err != nil {
		t.Fatalf("Info: %v", err)
	}
	if info.Name != "tushare" {
		t.Errorf("Name = %q, want tushare", info.Name)
	}
	if !strings.Contains(info.Desc, "Tushare") {
		t.Errorf("Desc = %q, want to mention Tushare", info.Desc)
	}
}