zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-06-29 15:31:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Go: fix test and remove unused code (#15909)

Browse Source 
### What problem does this PR solve?

1. Fix go test, some cases still failed.
2. Remove unused code.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
Jin Hai
2026-06-10 20:38:43 +08:00
committed by GitHub
parent e132173d1a
commit 3e4fb8cf1c
8 changed files with 17 additions and 360 deletions
Download Patch File Download Diff File Expand all files Collapse all files

235
internal/cli/admin_variables_test.go
View File

@@ -1,235 +0,0 @@
//
// 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 cli
import (
"encoding/json"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strconv"
"testing"
)
func TestParseAdminVariableCommands(t *testing.T) {
tests := []struct {
name string
input string
command string
varName string
varValue string
hasValue bool
adminMode bool
}{
{
name: "list variables",
input: "list vars;",
command: "list_variables",
adminMode: true,
},
{
name: "show variables by prefix",
input: "show var mail;",
command: "show_variable",
varName: "mail",
adminMode: true,
},
{
name: "set integer variable",
input: "set var mail.port 15;",
command: "set_variable",
varName: "mail.port",
varValue: "15",
hasValue: true,
adminMode: true,
},
{
name: "set quoted string variable",
input: `set var mail.server "local host";`,
command: "set_variable",
varName: "mail.server",
varValue: "local host",
hasValue: true,
adminMode: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
cmd, err := NewParser(tt.input).Parse(tt.adminMode)
if err != nil {
t.Fatalf("Parse() error = %v", err)
}
if cmd.Type != tt.command {
t.Fatalf("command type = %q, want %q", cmd.Type, tt.command)
}
if tt.varName != "" && cmd.Params["var_name"] != tt.varName {
t.Fatalf("var_name = %v, want %q", cmd.Params["var_name"], tt.varName)
}
if tt.hasValue && cmd.Params["var_value"] != tt.varValue {
t.Fatalf("var_value = %v, want %q", cmd.Params["var_value"], tt.varValue)
}
})
}
}
func newAdminTestClient(t *testing.T, handler http.HandlerFunc) (*RAGFlowClient, func()) {
t.Helper()
server := httptest.NewServer(handler)
serverURL, err := url.Parse(server.URL)
if err != nil {
t.Fatalf("parse test server URL: %v", err)
}
host, portText, err := net.SplitHostPort(serverURL.Host)
if err != nil {
t.Fatalf("split host port: %v", err)
}
port, err := strconv.Atoi(portText)
if err != nil {
t.Fatalf("parse port: %v", err)
}
client := NewRAGFlowClient("admin")
client.HTTPClient.Host = host
client.HTTPClient.Port = port
client.HTTPClient.client = server.Client()
client.HTTPClient.LoginToken = "test-token"
return client, server.Close
}
func TestListVariablesUsesAdminVariablesEndpoint(t *testing.T) {
client, closeServer := newAdminTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %s, want GET", r.Method)
return
}
if r.URL.Path != "/api/v1/admin/variables" {
t.Errorf("path = %s, want /api/v1/admin/variables", r.URL.Path)
return
}
if r.Header.Get("Authorization") != "test-token" {
t.Errorf("Authorization header = %q", r.Header.Get("Authorization"))
return
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"code": 0,
"message": "",
"data": []map[string]interface{}{
{"data_type": "string", "name": "mail.server", "source": "variable", "value": "localhost"},
},
})
})
defer closeServer()
resp, err := client.ListVariables(NewCommand("list_variables"))
if err != nil {
t.Fatalf("ListVariables() error = %v", err)
}
result := resp.(*CommonResponse)
if got := result.Data[0]["setting_type"]; got != "config" {
t.Fatalf("setting_type = %v, want config", got)
}
if _, ok := result.Data[0]["source"]; ok {
t.Fatalf("source column should be normalized away: %#v", result.Data[0])
}
}
func TestShowVariableSendsRequestedName(t *testing.T) {
client, closeServer := newAdminTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
t.Errorf("method = %s, want GET", r.Method)
return
}
body, err := io.ReadAll(r.Body)
if err != nil {
t.Errorf("read body: %v", err)
return
}
var request map[string]string
if err := json.Unmarshal(body, &request); err != nil {
t.Errorf("request body is not JSON: %v", err)
return
}
if request["var_name"] != "mail" {
t.Errorf("var_name = %q, want mail", request["var_name"])
return
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"code": 0,
"message": "",
"data": []map[string]interface{}{
{"data_type": "string", "name": "mail.server", "setting_type": "config", "value": "localhost"},
},
})
})
defer closeServer()
cmd := NewCommand("show_variable")
cmd.Params["var_name"] = "mail"
resp, err := client.ShowVariable(cmd)
if err != nil {
t.Fatalf("ShowVariable() error = %v", err)
}
result := resp.(*CommonResponse)
if got := result.Data[0]["name"]; got != "mail.server" {
t.Fatalf("name = %v, want mail.server", got)
}
}
func TestSetVariableReturnsServerConfirmation(t *testing.T) {
client, closeServer := newAdminTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPut {
t.Errorf("method = %s, want PUT", r.Method)
return
}
var request map[string]string
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Errorf("request body is not JSON: %v", err)
return
}
if request["var_name"] != "mail.server" {
t.Errorf("var_name = %q, want mail.server", request["var_name"])
return
}
if request["var_value"] != "localhost" {
t.Errorf("var_value = %q, want localhost", request["var_value"])
return
}
_ = json.NewEncoder(w).Encode(map[string]interface{}{
"code": 0,
"message": "Set variable successfully",
"data": nil,
})
})
defer closeServer()
cmd := NewCommand("set_variable")
cmd.Params["var_name"] = "mail.server"
cmd.Params["var_value"] = "localhost"
resp, err := client.SetVariable(cmd)
if err != nil {
t.Fatalf("SetVariable() error = %v", err)
}
result := resp.(*MessageResponse)
if result.Message != "Set variable successfully" {
t.Fatalf("message = %q, want Set variable successfully", result.Message)
}
}

23
internal/cli/cli.go
View File

@@ -422,24 +422,6 @@ func parseHostPort(hostPort string) (string, int, error) {
return host, port, nil
}
// looksLikeSQL checks if a string looks like a SQL command
func looksLikeSQL(s string) bool {
s = strings.ToUpper(strings.TrimSpace(s))
sqlPrefixes := []string{
"LIST ", "SHOW ", "CREATE ", "DROP ", "ALTER ",
"LOGIN ", "REGISTER ", "PING", "GRANT ", "REVOKE ",
"SET ", "UNSET ", "UPDATE ", "DELETE ", "INSERT ",
"SELECT ", "DESCRIBE ", "EXPLAIN ", "ADD ", "ENABLE ", "DISABLE ", "CHAT ", "USE", "THINK",
"REMOVE ",
}
for _, prefix := range sqlPrefixes {
if strings.HasPrefix(s, prefix) {
return true
}
}
return false
}
// PrintUsage prints the CLI usage information
func PrintUsage() {
fmt.Println(`RAGFlow CLI Client
@@ -510,11 +492,6 @@ type CLI struct {
Config *CommandLineConfig
}
// NewCLI creates a new CLI instance
//func NewCLI() (*CLI, error) {
// return NewCLIWithArgs(nil)
//}
func NewCLIWithConfig(commandLineConfig *CommandLineConfig) (*CLI, error) {
// Create liner first
line := liner.NewLiner()

2
internal/cli/common_command.go
View File

@@ -139,7 +139,7 @@ func (c *CLI) PingServer(iterations int) (ResponseIf, error) {
switch c.Config.CLIMode {
case AdminMode:
if err = json.Unmarshal(resp.Body, &result); err != nil {
return nil, fmt.Errorf("list users failed: invalid JSON (%w)", err)
return nil, fmt.Errorf("ping failed: invalid JSON (%w)", err)
}
case APIMode:
if string(resp.Body) == "pong" {

80
internal/proto/ingestion.proto
View File

@@ -1,80 +0,0 @@
syntax = "proto3";
package common;
option go_package = "./;common";
service IngestionManager {
rpc Action(stream IngestionMessage) returns (stream AdminMessage);
}
message IngestionMessage {
string ingestor_id = 1;
string message_type = 2; // REGISTER, HEARTBEAT, TASK_RESULT, TASK_PROGRESS, PULL_REQUEST
RegisterInfo register_info = 3;
HeartbeatInfo heartbeat_info = 4;
TaskResult task_result = 5;
TaskProgress task_progress = 6;
}
message AdminMessage {
string message_type = 1; // TASK_ASSIGNMENT, ACK, PONG, RECONNECT
TaskAssignment task_assignment = 2;
AckInfo ack_info = 3;
string error_message = 4;
}
message RegisterInfo {
int32 max_concurrency = 1;
repeated string supported_doc_types = 2;
string version = 3;
string name = 4;
}
message HeartbeatInfo {
repeated TaskState task_states = 1;
repeated string delete_task_ids = 2;
float cpu_usage = 3; // percentage
float vms_usage = 4; // absolute value
float rss_usage = 5; // absolute value
int64 process_id = 6; // pid
}
message TaskState {
string task_id = 1;
string status = 2; // PENDING, RUNNING, COMPLETED, FAILED, CANCELLED
string error_message = 3;
int64 estimated_remaining_time_seconds = 4;
string come_from = 5;
int64 start_time = 6;
}
message TaskAssignment {
string task_id = 1;
string task_type = 2;
string config = 3;
string come_from = 4;
string assigned_to = 5;
}
message TaskResult {
string task_id = 1;
string status = 2; // COMPLETED, FAILED, CANCELLED
string error_message = 3;
}
message TaskProgress {
string task_id = 1;
int32 progress = 2;
string info = 3;
}
message AckInfo {
string task_id = 1;
bool success = 2;
string message = 3;
}

7
internal/service/mcp.go
View File

@@ -28,7 +28,6 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/mcpclient"
"ragflow/internal/utility"
"gorm.io/gorm"
@@ -491,7 +490,7 @@ func (s *MCPService) ImportServers(tenantID string, servers map[string]map[strin
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
tools, fetchErr := mcpclient.FetchTools(ctx, mcpclient.FetchOptions{
tools, fetchErr := utility.FetchTools(ctx, utility.FetchOptions{
URL: url,
ServerType: stype,
Headers: headers,
@@ -597,7 +596,7 @@ func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[str
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
tools, err := mcpclient.FetchTools(ctx, mcpclient.FetchOptions{
tools, err := utility.FetchTools(ctx, utility.FetchOptions{
URL: req.URL,
ServerType: req.ServerType,
Headers: headers,
@@ -628,7 +627,7 @@ func (s *MCPService) TestServer(mcpID string, req *TestServerRequest) ([]map[str
// toolsAsMap mirrors Python's `{tool["name"]: tool ...}` shape used when
// persisting variables.tools.
func toolsAsMap(tools []mcpclient.Tool) map[string]interface{} {
func toolsAsMap(tools []utility.Tool) map[string]interface{} {
m := map[string]interface{}{}
for _, t := range tools {
if t.Raw != nil {

12
internal/mcpclient/client.go → internal/utility/mcp_client.go
View File

@@ -28,7 +28,7 @@
//
// The full Python implementation lives in common/mcp_tool_call_conn.py; this
// is a reduced port focused on tools/list discovery.
package mcpclient
package utility
import (
"bufio"
@@ -43,8 +43,6 @@ import (
"strings"
"sync"
"time"
"ragflow/internal/utility"
)
// Transport identifiers. Mirrors common.constants.MCPServerType.
@@ -94,14 +92,14 @@ func FetchTools(ctx context.Context, opts FetchOptions) ([]Tool, error) {
opts.Timeout = 10 * time.Second
}
hostname, resolvedIP, err := utility.AssertURLSafe(opts.URL)
hostname, resolvedIP, err := AssertURLSafe(opts.URL)
if err != nil {
return nil, err
}
opts.pinHostname = hostname
opts.pinIP = resolvedIP
if opts.HTTPClient == nil {
opts.HTTPClient = utility.PinnedHTTPClient(hostname, resolvedIP, opts.Timeout)
opts.HTTPClient = PinnedHTTPClient(hostname, resolvedIP, opts.Timeout)
}
headers, headerErr := renderHeaders(opts.Headers, opts.Variables)
@@ -346,11 +344,11 @@ func fetchToolsSSE(ctx context.Context, endpoint string, headers map[string]stri
// differs from the original SSE host — swap in a fresh pinned
// client so the dial-time IP override still applies.
postClient := client
if postHost, postIP, vErr := utility.AssertURLSafe(postURL); vErr != nil {
if postHost, postIP, vErr := AssertURLSafe(postURL); vErr != nil {
return nil, vErr
} else if u, perr := url.Parse(postURL); perr == nil && u.Hostname() != "" {
if u.Hostname() != originalHost(endpoint) {
postClient = utility.PinnedHTTPClient(postHost, postIP, sseTimeoutFrom(ctx))
postClient = PinnedHTTPClient(postHost, postIP, sseTimeoutFrom(ctx))
}
}

10
internal/mcpclient/client_test.go → internal/utility/mcp_client_test.go
View File

@@ -14,7 +14,7 @@
// limitations under the License.
//
package mcpclient
package utility
import (
"context"
@@ -27,21 +27,19 @@ import (
"sync/atomic"
"testing"
"time"
"ragflow/internal/utility"
)
// allowLoopbackForTests overrides the SSRF guard's resolver so 127.0.0.1
// targets used by httptest are accepted by AssertURLSafe.
func allowLoopbackForTests(t *testing.T) func() {
t.Helper()
orig := utility.LookupHost
utility.LookupHost = func(host string) ([]string, error) {
orig := LookupHost
LookupHost = func(host string) ([]string, error) {
// Return a public IPv4 so the guard sees the host as global; the
// httptest server is on loopback but we connect via raw URL.
return []string{"8.8.8.8"}, nil
}
return func() { utility.LookupHost = orig }
return func() { LookupHost = orig }
}
func TestFetchToolsStreamableHTTPJSON(t *testing.T) {

8
run_go_tests.sh
View File

@@ -9,15 +9,15 @@ PACKAGES=(
"./internal/common/..."
"./internal/dao/..."
"./internal/engine/..."
"./internal/entity/..."
"./internal/handler/..."
"./internal/logger/..."
"./internal/model/..."
"./internal/ingestion/..."
"./internal/router/..."
"./internal/server/..."
# "./internal/service/..."
"./internal/service/..."
"./internal/storage/..."
"./internal/tokenizer/..."
# "./internal/utility/..."
"./internal/utility/..."
)
echo "Running tests for specific packages..."