FranklinWH-Toolbox/service/cmd/mcp/tools/middleware/logging_test.go

package middleware

import (
	"context"
	"strings"
	"testing"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

// setupObservedLogger 替换全局 logger 为可观测 logger，返回观测器与还原函数。
func setupObservedLogger(t *testing.T) (*observer.ObservedLogs, func()) {
	t.Helper()
	core, obs := observer.New(zapcore.InfoLevel)
	orig := zap.L()
	zap.ReplaceGlobals(zap.New(core))
	return obs, func() { zap.ReplaceGlobals(orig) }
}

func newReq(name string, args map[string]any) mcp.CallToolRequest {
	var req mcp.CallToolRequest
	req.Params.Name = name
	req.Params.Arguments = args
	return req
}

func wrap(next server.ToolHandlerFunc) server.ToolHandlerFunc {
	return LoggingMiddleware()(next)
}

func TestLogging_Success(t *testing.T) {
	obs, restore := setupObservedLogger(t)
	defer restore()

	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return mcp.NewToolResultText(`{"ok":true}`), nil
	})
	_, err := h(context.Background(), newReq("service_site_query_info", map[string]any{"sn": "SN123"}))
	if err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	logs := obs.All()
	if len(logs) != 2 {
		t.Fatalf("want 2 log entries (start+success), got %d: %+v", len(logs), logs)
	}
	if logs[0].Message != "[tool] call start" {
		t.Errorf("first log = %q, want start", logs[0].Message)
	}
	if logs[1].Message != "[tool] call success" || logs[1].Level != zapcore.InfoLevel {
		t.Errorf("second log = %q level=%v, want success/info", logs[1].Message, logs[1].Level)
	}
	// status 字段
	if got := logs[1].ContextMap()["status"]; got != "success" {
		t.Errorf("status = %v, want success", got)
	}
}

func TestLogging_BusinessError(t *testing.T) {
	obs, restore := setupObservedLogger(t)
	defer restore()

	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return mcp.NewToolResultError("查询站点信息失败: not found"), nil
	})
	if _, err := h(context.Background(), newReq("service_site_query_info", nil)); err != nil {
		t.Fatalf("unexpected err: %v", err)
	}

	last := obs.All()[len(obs.All())-1]
	if last.Message != "[tool] call business_error" || last.Level != zapcore.WarnLevel {
		t.Fatalf("last log = %q level=%v, want business_error/warn", last.Message, last.Level)
	}
	if got := last.ContextMap()["status"]; got != "business_error" {
		t.Errorf("status = %v, want business_error", got)
	}
}

func TestLogging_GoError(t *testing.T) {
	obs, restore := setupObservedLogger(t)
	defer restore()

	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return nil, context.DeadlineExceeded
	})
	if _, err := h(context.Background(), newReq("data_mongo_query", nil)); err == nil {
		t.Fatal("expected error to propagate")
	}

	last := obs.All()[len(obs.All())-1]
	if last.Message != "[tool] call error" || last.Level != zapcore.ErrorLevel {
		t.Fatalf("last log = %q level=%v, want error/error", last.Message, last.Level)
	}
}

func TestLogging_Panic(t *testing.T) {
	obs, restore := setupObservedLogger(t)
	defer restore()

	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		panic("boom")
	})

	defer func() {
		if r := recover(); r == nil {
			t.Fatal("expected panic to be re-thrown")
		}
		last := obs.All()[len(obs.All())-1]
		if last.Message != "[tool] call panic" || last.Level != zapcore.ErrorLevel {
			t.Fatalf("last log = %q level=%v, want panic/error", last.Message, last.Level)
		}
		if _, ok := last.ContextMap()["stack"]; !ok {
			t.Error("panic log missing stack field")
		}
	}()

	_, _ = h(context.Background(), newReq("service_site_query_info", nil))
}

// TestLogging_ParamMasking 验证参数中的敏感字段被脱敏后再写日志。
func TestLogging_ParamMasking(t *testing.T) {
	obs, restore := setupObservedLogger(t)
	defer restore()

	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return mcp.NewToolResultText("ok"), nil
	})
	// phone 是常见 PII，MaskJSONBytes 应对其脱敏
	_, _ = h(context.Background(), newReq("service_site_query_info", map[string]any{"phone": "13800138000"}))

	start := obs.All()[0]
	params, _ := start.ContextMap()["params"].(string)
	if strings.Contains(params, "13800138000") {
		t.Errorf("params log leaked raw phone: %q", params)
	}
}