// Copyright (c) 2026 Lark Technologies Pte. Ltd. // SPDX-License-Identifier: MIT package common import ( "context" "errors" "io" "net/http" "strings" "testing" "github.com/larksuite/cli/errs" "github.com/larksuite/cli/internal/output" ) type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } func TestDoMCPCallUnauthorizedHTTPError(t *testing.T) { t.Parallel() client := &http.Client{ Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { return &http.Response{ StatusCode: http.StatusUnauthorized, Status: "401 Unauthorized", Body: io.NopCloser(strings.NewReader("unauthorized")), }, nil }), } _, err := DoMCPCall(context.Background(), client, "fetch-doc", map[string]interface{}{"doc_id": "doc_1"}, "uat-token", "https://example.com/mcp", false) if got := output.ExitCodeOf(err); got != output.ExitAuth { t.Fatalf("expected auth exit code (%d), got %d", output.ExitAuth, got) } } func TestDoMCPCallJSONRPCErrorUsesLarkClassification(t *testing.T) { t.Parallel() client := &http.Client{ Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { return &http.Response{ StatusCode: http.StatusOK, Status: "200 OK", Body: io.NopCloser(strings.NewReader(`{"error":{"code":99991668,"message":"user_access_token invalid"}}`)), }, nil }), } _, err := DoMCPCall(context.Background(), client, "fetch-doc", map[string]interface{}{"doc_id": "doc_1"}, "uat-token", "https://example.com/mcp", false) if err == nil { t.Fatal("expected error, got nil") } if got := output.ExitCodeOf(err); got != output.ExitAuth { t.Fatalf("expected auth exit code (%d), got %d", output.ExitAuth, got) } var authErr *errs.AuthenticationError if !errors.As(err, &authErr) { t.Fatalf("expected *errs.AuthenticationError, got %T: %v", err, err) } } func TestDoMCPCallSetsHeadersAndUnwrapsResult(t *testing.T) { t.Parallel() var seen *http.Request client := &http.Client{ Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { seen = req return &http.Response{ StatusCode: http.StatusOK, Status: "200 OK", Body: io.NopCloser(strings.NewReader(`{"result":{"jsonrpc":"2.0","result":{"ok":true}}}`)), }, nil }), } got, err := DoMCPCall(context.Background(), client, "fetch-doc", map[string]interface{}{"doc_id": "doc_1"}, "tat-token", "https://example.com/mcp", true) if err != nil { t.Fatalf("unexpected error: %v", err) } result, ok := got.(map[string]interface{}) if !ok || result["ok"] != true { t.Fatalf("unexpected result: %#v", got) } if seen == nil { t.Fatalf("expected request to be captured") } if seen.Header.Get("X-Lark-MCP-TAT") != "tat-token" { t.Fatalf("expected bot token header, got %q", seen.Header.Get("X-Lark-MCP-TAT")) } if seen.Header.Get("X-Lark-MCP-Allowed-Tools") != "fetch-doc" { t.Fatalf("expected allowed tools header, got %q", seen.Header.Get("X-Lark-MCP-Allowed-Tools")) } } func TestNormalizeMCPToolResult(t *testing.T) { t.Parallel() tests := []struct { name string raw interface{} wantKey string wantVal interface{} wantErr string }{ { name: "map result", raw: map[string]interface{}{"ok": true}, wantKey: "ok", wantVal: true, }, { name: "text result", raw: "plain text", wantKey: "message", wantVal: "plain text", }, { name: "scalar result", raw: 42, wantKey: "result", wantVal: 42, }, { name: "map error field", raw: map[string]interface{}{"error": "permission denied"}, wantErr: "MCP: permission denied", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() got, err := normalizeMCPToolResult(tt.raw) if tt.wantErr != "" { if err == nil || !strings.Contains(err.Error(), tt.wantErr) { t.Fatalf("expected error containing %q, got %v", tt.wantErr, err) } return } if err != nil { t.Fatalf("unexpected error: %v", err) } if got[tt.wantKey] != tt.wantVal { t.Fatalf("unexpected result: %#v", got) } }) } } func TestExtractMCPResult(t *testing.T) { t.Parallel() jsonResult := ExtractMCPResult(map[string]interface{}{ "content": []interface{}{ map[string]interface{}{ "type": "text", "text": `{"doc_id":"doc_1"}`, }, }, }) resultMap, ok := jsonResult.(map[string]interface{}) if !ok || resultMap["doc_id"] != "doc_1" { t.Fatalf("unexpected parsed json result: %#v", jsonResult) } textResult := ExtractMCPResult(map[string]interface{}{ "content": []interface{}{ map[string]interface{}{"type": "text", "text": "line1"}, map[string]interface{}{"type": "text", "text": "line2"}, }, }) if textResult != "line1\nline2" { t.Fatalf("unexpected text result: %#v", textResult) } }