From e46ef82aac296c4efef0ab3351b29f369a9d8cc9 Mon Sep 17 00:00:00 2001 From: Zhang Rui Date: Thu, 19 Mar 2026 22:04:13 +0800 Subject: [PATCH] feat(upload): implement WeCom media upload via WebSocket API --- pkg/agent/loop.go | 2 + pkg/agent/loop_mcp.go | 2 +- pkg/bus/types.go | 1 + pkg/channels/wecom/aibot_ws.go | 85 +++---- pkg/channels/wecom/aibot_ws_upload.go | 234 ++++++++++++++++++ pkg/channels/wecom/aibot_ws_upload_test.go | 268 +++++++++++++++++++++ 6 files changed, 550 insertions(+), 42 deletions(-) create mode 100644 pkg/channels/wecom/aibot_ws_upload.go create mode 100644 pkg/channels/wecom/aibot_ws_upload_test.go diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index edb0994c2..6501ef7c7 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -193,6 +193,7 @@ func registerSharedTools( Channel: channel, ChatID: chatID, Content: content, + Final: true, }) }) agent.Tools.Register(messageTool) @@ -315,6 +316,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { Channel: msg.Channel, ChatID: msg.ChatID, Content: response, + Final: true, }) logger.InfoCF("agent", "Published outbound response", map[string]any{ diff --git a/pkg/agent/loop_mcp.go b/pkg/agent/loop_mcp.go index 97debbc33..e0840b7d3 100644 --- a/pkg/agent/loop_mcp.go +++ b/pkg/agent/loop_mcp.go @@ -64,7 +64,7 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error { return nil } - if al.cfg.Tools.MCP.Servers == nil || len(al.cfg.Tools.MCP.Servers) == 0 { + if len(al.cfg.Tools.MCP.Servers) == 0 { logger.WarnCF("agent", "MCP is enabled but no servers are configured, skipping MCP initialization", nil) return nil } diff --git a/pkg/bus/types.go b/pkg/bus/types.go index 12da3f1dd..f95c9ad09 100644 --- a/pkg/bus/types.go +++ b/pkg/bus/types.go @@ -34,6 +34,7 @@ type OutboundMessage struct { ChatID string `json:"chat_id"` Content string `json:"content"` ReplyToMessageID string `json:"reply_to_message_id,omitempty"` + Final bool `json:"final,omitempty"` } // MediaPart describes a single media attachment to send. diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go index 830e763b9..e75ff9e50 100644 --- a/pkg/channels/wecom/aibot_ws.go +++ b/pkg/channels/wecom/aibot_ws.go @@ -38,10 +38,8 @@ const ( wsInitialReconnect = time.Second // WeCom requires finish=true within 6 minutes of the first stream frame. - // wsStreamTickInterval controls how often we send an in-progress hint. // wsStreamMaxDuration is a safety margin below the 6-minute hard limit. - wsStreamTickInterval = 30 * time.Second - wsStreamMaxDuration = 5*time.Minute + 30*time.Second + wsStreamMaxDuration = 5*time.Minute + 30*time.Second // wsImageDownloadTimeout caps the time we spend downloading an inbound image. wsImageDownloadTimeout = 30 * time.Second @@ -286,6 +284,11 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) // and there will be no matching entry in reqStates; fall through to proactive push. task, route, ok := c.getReqState(msg.ChatID) if !ok { + if !msg.Final { + // Intermediate message (tool feedback, etc.) with no active WS session. + // Nothing to stream to; silently drop it. + return nil + } // No req_id record found — this is a cron/scheduler-originated message. // Send it as a proactive markdown push using the chat ID directly. logger.InfoCF("wecom_aibot", "Send: no req_id state, delivering via proactive push (cron/scheduler)", @@ -299,6 +302,12 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) } if task == nil { + if !msg.Final { + // Intermediate message: stream for this req_id has already closed (deadline + // or new-message cancel). Nothing to do. + return nil + } + if time.Now().Before(route.ReadyAt) { // Keep using aibot_respond_msg within stream window; do not proactively // push unless wsStreamMaxDuration has elapsed. @@ -318,6 +327,19 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage) return nil } + if !msg.Final { + // Intermediate message (tool feedback, context notices, etc.) with an active + // stream. Send directly as a non-finishing stream chunk so the user sees + // live progress without touching answerCh (which the goroutine reserves for + // the single final answer). + for _, chunk := range splitWSContent(msg.Content, wsStreamMaxContentBytes) { + c.wsSendStreamChunk(task.ReqID, task.StreamID, false, chunk) + } + return nil + } + + // Final answer: deliver to the waiting goroutine via answerCh, which will + // send it with finish=true and clean up the stream state. // Non-blocking fast path: when answerCh has space, deliver without racing // against task.ctx.Done() (which fires when the task is canceled by a new // incoming message, but the response must still be sent). @@ -852,37 +874,34 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask( c.HandleMessage(taskCtx, peer, reqID, userID, reqID, content, mediaRefs, metadata, sender) - // Wait for the agent reply. While waiting, send periodic finish=false - // hints so the user knows processing is still in progress. + // Wait for the agent reply. // WeCom requires finish=true within 6 minutes of the first stream frame; // wsStreamMaxDuration enforces that limit with a safety margin. - waitHints := []string{ - "⏳ Processing, please wait...", - "⏳ Still processing, please wait...", - "⏳ Almost there, please wait...", - } - ticker := time.NewTicker(wsStreamTickInterval) - defer ticker.Stop() deadlineTimer := time.NewTimer(wsStreamMaxDuration) defer deadlineTimer.Stop() - tickCount := 0 for { select { case answer := <-task.answerCh: - // Split the answer into byte-bounded chunks and send as stream frames. - // All but the last carry finish=false; the final frame closes the stream. + + // Split the answer into byte-bounded chunks. + // Each chunk must fit within wsStreamMaxContentBytes (20480 bytes). + // WeCom's stream update semantics REPLACE content (not append), so all + // chunks beyond the first need a fresh stream.id — otherwise only the + // last chunk would be visible to the user. + // + // Delivery plan: + // chunk[0] — close the already-opened stream (finish=true, same streamID) + // chunk[1+] — create + immediately close a new stream per chunk chunks := splitWSContent(answer, wsStreamMaxContentBytes) for i, chunk := range chunks { - c.wsSendStreamChunk(reqID, streamID, i == len(chunks)-1, chunk) + sid := streamID + if i > 0 { + sid = wsGenerateID() + } + c.wsSendStreamChunk(reqID, sid, true, chunk) } c.deleteReqState(reqID) return - case <-ticker.C: - hint := waitHints[tickCount%len(waitHints)] - tickCount++ - logger.DebugCF("wecom_aibot", "Sending stream progress hint", - map[string]any{"chat_id": actualChatID, "tick": tickCount}) - c.wsSendStreamChunk(reqID, streamID, false, hint) case <-deadlineTimer.C: logger.WarnCF("wecom_aibot", "Stream response deadline reached, closing stream; late reply will be pushed", @@ -1028,26 +1047,10 @@ func (c *WeComAIBotWSChannel) wsSendActivePush(chatID string, chatType uint32, c } // writeWSAndWait writes cmd to the active connection and validates the command response. +// It is a thin wrapper around callWSCommand that discards the response body. func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error { - if cmd.Headers.ReqID == "" { - return fmt.Errorf("req_id is empty") - } - - c.connMu.Lock() - conn := c.conn - c.connMu.Unlock() - if conn == nil { - return fmt.Errorf("websocket not connected") - } - - resp, err := c.sendAndWait(conn, cmd.Headers.ReqID, cmd, timeout) - if err != nil { - return err - } - if resp.ErrCode != 0 { - return fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, resp.ErrCode, resp.ErrMsg) - } - return nil + _, err := c.callWSCommand(cmd, timeout) + return err } // cancelAllTasks cancels every pending agent task; called when the connection drops. diff --git a/pkg/channels/wecom/aibot_ws_upload.go b/pkg/channels/wecom/aibot_ws_upload.go new file mode 100644 index 000000000..8dab028cd --- /dev/null +++ b/pkg/channels/wecom/aibot_ws_upload.go @@ -0,0 +1,234 @@ +package wecom + +import ( + "context" + "crypto/md5" + "encoding/base64" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// Upload temporary media constants. +// Ref: https://developer.work.weixin.qq.com/document/path/101463#上传临时素材 +const ( + // wsUploadChunkSize is the maximum raw (pre-base64) bytes per chunk. + // The API limit is 512 KB before base64 encoding. + wsUploadChunkSize = 512 << 10 // 512 KB + + // wsUploadMaxChunks is the maximum number of chunks per upload session. + wsUploadMaxChunks = 100 + + wsUploadInitTimeout = 30 * time.Second + wsUploadChunkTimeout = 60 * time.Second // larger: base64 payload may be big + wsUploadFinishTimeout = 30 * time.Second +) + +// ---- Request / response body types ---- + +// wsUploadInitBody is the body for aibot_upload_media_init. +type wsUploadInitBody struct { + Type string `json:"type"` + Filename string `json:"filename"` + TotalSize int64 `json:"total_size"` + TotalChunks int `json:"total_chunks"` + MD5 string `json:"md5,omitempty"` +} + +// wsUploadInitResponse is the body received in the aibot_upload_media_init response. +type wsUploadInitResponse struct { + UploadID string `json:"upload_id"` +} + +// wsUploadChunkBody is the body for aibot_upload_media_chunk. +type wsUploadChunkBody struct { + UploadID string `json:"upload_id"` + ChunkIndex int `json:"chunk_index"` + Base64Data string `json:"base64_data"` +} + +// wsUploadFinishBody is the body for aibot_upload_media_finish. +type wsUploadFinishBody struct { + UploadID string `json:"upload_id"` +} + +// wsUploadFinishResponse is the body received in the aibot_upload_media_finish response. +type wsUploadFinishResponse struct { + Type string `json:"type"` + MediaID string `json:"media_id"` + CreatedAt string `json:"created_at"` +} + +// ---- Public API ---- + +// UploadWSMedia uploads a local file as a temporary WeCom media asset via the +// WebSocket long-connection API and returns the resulting media_id (valid 3 days). +// +// mediaType must be one of: +// - "image" — PNG / JPG(JPEG) / GIF, ≤ 2 MB +// - "voice" — AMR, ≤ 2 MB +// - "video" — MP4, ≤ 10 MB +// - "file" — any format, ≤ 20 MB +// +// filename is the logical file name sent to WeCom. When empty it is derived +// from the base name of filePath. +// +// The upload is split into at most wsUploadMaxChunks chunks of wsUploadChunkSize +// bytes each. The full operation must complete within wsUploadSessionTTL (30 min). +// Each individual command call is also independently time-bounded. +func (c *WeComAIBotWSChannel) UploadWSMedia( + ctx context.Context, + filePath, filename, mediaType string, +) (string, error) { + // ---- Validate media type ---- + switch mediaType { + case "image", "voice", "video", "file": + default: + return "", fmt.Errorf("unsupported media type %q: must be image, voice, video, or file", mediaType) + } + + // ---- Read file ---- + data, err := os.ReadFile(filePath) + if err != nil { + return "", fmt.Errorf("read file: %w", err) + } + totalSize := int64(len(data)) + if totalSize == 0 { + return "", fmt.Errorf("file is empty: %s", filePath) + } + + if filename == "" { + filename = filepath.Base(filePath) + } + + // ---- Compute chunk count ---- + totalChunks := int((totalSize + wsUploadChunkSize - 1) / wsUploadChunkSize) + if totalChunks > wsUploadMaxChunks { + return "", fmt.Errorf( + "file too large: requires %d chunks but limit is %d (max ~%d MB)", + totalChunks, wsUploadMaxChunks, + int64(wsUploadMaxChunks)*wsUploadChunkSize>>20, + ) + } + + // ---- Compute MD5 ---- + rawMD5 := md5.Sum(data) //nolint:gosec // MD5 is required by the WeCom API spec + fileMD5 := fmt.Sprintf("%x", rawMD5) + + // ---- Step 1: Init ---- + initEnv, err := c.callWSCommand(wsCommand{ + Cmd: "aibot_upload_media_init", + Headers: wsHeaders{ReqID: wsGenerateID()}, + Body: wsUploadInitBody{ + Type: mediaType, + Filename: filename, + TotalSize: totalSize, + TotalChunks: totalChunks, + MD5: fileMD5, + }, + }, wsUploadInitTimeout) + if err != nil { + return "", fmt.Errorf("upload init: %w", err) + } + var initResp wsUploadInitResponse + if err = json.Unmarshal(initEnv.Body, &initResp); err != nil { + return "", fmt.Errorf("parse upload init response: %w", err) + } + if initResp.UploadID == "" { + return "", fmt.Errorf("upload init returned empty upload_id") + } + uploadID := initResp.UploadID + + logger.InfoCF("wecom_aibot", "Media upload initialized", map[string]any{ + "upload_id": uploadID, + "type": mediaType, + "filename": filename, + "total_size": totalSize, + "total_chunks": totalChunks, + }) + + // ---- Step 2: Upload chunks ---- + for i := 0; i < totalChunks; i++ { + if ctxErr := ctx.Err(); ctxErr != nil { + return "", fmt.Errorf("upload aborted before chunk %d: %w", i, ctxErr) + } + + start := int64(i) * wsUploadChunkSize + end := start + wsUploadChunkSize + if end > totalSize { + end = totalSize + } + + _, err = c.callWSCommand(wsCommand{ + Cmd: "aibot_upload_media_chunk", + Headers: wsHeaders{ReqID: wsGenerateID()}, + Body: wsUploadChunkBody{ + UploadID: uploadID, + ChunkIndex: i, + Base64Data: base64.StdEncoding.EncodeToString(data[start:end]), + }, + }, wsUploadChunkTimeout) + if err != nil { + return "", fmt.Errorf("upload chunk %d/%d: %w", i, totalChunks-1, err) + } + + logger.DebugCF("wecom_aibot", "Media chunk uploaded", map[string]any{ + "upload_id": uploadID, + "chunk": i, + "total": totalChunks, + }) + } + + // ---- Step 3: Finish ---- + finishEnv, err := c.callWSCommand(wsCommand{ + Cmd: "aibot_upload_media_finish", + Headers: wsHeaders{ReqID: wsGenerateID()}, + Body: wsUploadFinishBody{UploadID: uploadID}, + }, wsUploadFinishTimeout) + if err != nil { + return "", fmt.Errorf("upload finish: %w", err) + } + var finishResp wsUploadFinishResponse + if err := json.Unmarshal(finishEnv.Body, &finishResp); err != nil { + return "", fmt.Errorf("parse upload finish response: %w", err) + } + if finishResp.MediaID == "" { + return "", fmt.Errorf("upload finish returned empty media_id") + } + + logger.InfoCF("wecom_aibot", "Media upload complete", map[string]any{ + "upload_id": uploadID, + "media_id": finishResp.MediaID, + "type": finishResp.Type, + }) + return finishResp.MediaID, nil +} + +// ---- Internal helper ---- + +// callWSCommand sends a WebSocket command and returns the raw server envelope. +// It validates the errcode field and returns an error when non-zero. +// Use callWSCommand (over writeWSAndWait) when the response body must be inspected. +func (c *WeComAIBotWSChannel) callWSCommand(cmd wsCommand, timeout time.Duration) (wsEnvelope, error) { + if cmd.Headers.ReqID == "" { + return wsEnvelope{}, fmt.Errorf("req_id is empty") + } + c.connMu.Lock() + conn := c.conn + c.connMu.Unlock() + if conn == nil { + return wsEnvelope{}, fmt.Errorf("websocket not connected") + } + env, err := c.sendAndWait(conn, cmd.Headers.ReqID, cmd, timeout) + if err != nil { + return wsEnvelope{}, err + } + if env.ErrCode != 0 { + return wsEnvelope{}, fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, env.ErrCode, env.ErrMsg) + } + return env, nil +} diff --git a/pkg/channels/wecom/aibot_ws_upload_test.go b/pkg/channels/wecom/aibot_ws_upload_test.go new file mode 100644 index 000000000..90f1961e0 --- /dev/null +++ b/pkg/channels/wecom/aibot_ws_upload_test.go @@ -0,0 +1,268 @@ +package wecom + +import ( + "context" + "encoding/base64" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +// --------------------------------------------------------------------------- +// UploadWSMedia — pre-flight validation tests (no real WS required) +// --------------------------------------------------------------------------- + +// TestUploadWSMedia_InvalidMediaType verifies that unsupported media types are +// rejected before any network activity occurs. +func TestUploadWSMedia_InvalidMediaType(t *testing.T) { + ch := newTestWSChannel(t) + + for _, bad := range []string{"", "pdf", "audio", "IMAGE", "gif"} { + _, err := ch.UploadWSMedia(context.Background(), "/dev/null", "test.pdf", bad) + if err == nil { + t.Errorf("expected error for media type %q, got nil", bad) + } + } +} + +// TestUploadWSMedia_FileNotFound verifies that a non-existent path returns an error. +func TestUploadWSMedia_FileNotFound(t *testing.T) { + ch := newTestWSChannel(t) + _, err := ch.UploadWSMedia(context.Background(), "/nonexistent/path/to/file.bin", "", "file") + if err == nil { + t.Fatal("expected error for missing file, got nil") + } +} + +// TestUploadWSMedia_EmptyFile verifies that an empty file is rejected before +// the upload starts. +func TestUploadWSMedia_EmptyFile(t *testing.T) { + ch := newTestWSChannel(t) + + tmp, err := os.CreateTemp(t.TempDir(), "empty-*.bin") + if err != nil { + t.Fatal(err) + } + tmp.Close() + + _, err = ch.UploadWSMedia(context.Background(), tmp.Name(), "", "file") + if err == nil { + t.Fatal("expected error for empty file, got nil") + } +} + +// TestUploadWSMedia_FileTooLarge verifies that a file requiring more than +// wsUploadMaxChunks chunks is rejected before the upload starts. +func TestUploadWSMedia_FileTooLarge(t *testing.T) { + ch := newTestWSChannel(t) + + // Synthesize a file that would need 101 chunks (one byte over the limit). + size := int64(wsUploadMaxChunks)*int64(wsUploadChunkSize) + 1 + + tmp, err := os.CreateTemp(t.TempDir(), "large-*.bin") + if err != nil { + t.Fatal(err) + } + truncErr := tmp.Truncate(size) + tmp.Close() + if truncErr != nil { + t.Fatal(truncErr) + } + + _, err = ch.UploadWSMedia(context.Background(), tmp.Name(), "", "file") + if err == nil { + t.Fatal("expected error for oversized file, got nil") + } +} + +// TestUploadWSMedia_NotConnected verifies that calling UploadWSMedia when no +// WebSocket connection is active returns a clear error. +func TestUploadWSMedia_NotConnected(t *testing.T) { + ch := newTestWSChannel(t) + + tmp := filepath.Join(t.TempDir(), "test.txt") + if err := os.WriteFile(tmp, []byte("hello world"), 0o600); err != nil { + t.Fatal(err) + } + + // conn is nil by default in a test channel — callWSCommand must return an error. + _, err := ch.UploadWSMedia(context.Background(), tmp, "", "file") + if err == nil { + t.Fatal("expected error when conn is nil, got nil") + } +} + +// TestUploadWSMedia_ContextCanceledPreFlight verifies that a pre-canceled +// context is detected during chunk iteration (even before any WS call, when +// the connection is nil the WS error takes precedence for a single-chunk file, +// so we test this path with a file that triggers ctx.Err() first). +func TestUploadWSMedia_ContextCanceledPreFlight(t *testing.T) { + ch := newTestWSChannel(t) + + tmp := filepath.Join(t.TempDir(), "test.txt") + if err := os.WriteFile(tmp, []byte("hello"), 0o600); err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // pre-cancel before the call + + // With conn == nil, callWSCommand will return "websocket not connected". + // The pre-canceled context is also an error; either one satisfies the test. + _, err := ch.UploadWSMedia(ctx, tmp, "", "file") + if err == nil { + t.Fatal("expected error for canceled context or nil conn, got nil") + } +} + +// --------------------------------------------------------------------------- +// callWSCommand — unit tests without a live connection +// --------------------------------------------------------------------------- + +// TestCallWSCommand_EmptyReqID verifies that callWSCommand rejects a command +// whose req_id is empty before touching the connection. +func TestCallWSCommand_EmptyReqID(t *testing.T) { + ch := newTestWSChannel(t) + cmd := wsCommand{ + Cmd: "aibot_upload_media_init", + Body: map[string]string{"type": "file"}, + // Headers.ReqID intentionally left empty + } + _, err := ch.callWSCommand(cmd, 5*time.Second) + if err == nil { + t.Fatal("expected error for empty req_id, got nil") + } +} + +// TestCallWSCommand_NotConnected verifies that callWSCommand returns a clear +// error when the WebSocket connection is nil. +func TestCallWSCommand_NotConnected(t *testing.T) { + ch := newTestWSChannel(t) + cmd := wsCommand{ + Cmd: "aibot_upload_media_init", + Headers: wsHeaders{ReqID: "test-req-id"}, + Body: map[string]string{"type": "file"}, + } + _, err := ch.callWSCommand(cmd, 5*time.Second) + if err == nil { + t.Fatal("expected error when conn is nil, got nil") + } +} + +// --------------------------------------------------------------------------- +// Upload body / response type round-trip tests +// --------------------------------------------------------------------------- + +// TestWsUploadChunkCount verifies the integer chunk-count formula used inside +// UploadWSMedia without needing the math package. +func TestWsUploadChunkCount(t *testing.T) { + tests := []struct { + name string + size int64 + wantChunks int + }{ + {"one byte", 1, 1}, + {"exactly one chunk", int64(wsUploadChunkSize), 1}, + {"one byte over", int64(wsUploadChunkSize) + 1, 2}, + {"max allowed", int64(wsUploadMaxChunks) * int64(wsUploadChunkSize), wsUploadMaxChunks}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := int((tc.size + int64(wsUploadChunkSize) - 1) / int64(wsUploadChunkSize)) + if got != tc.wantChunks { + t.Errorf("size=%d: got %d chunks, want %d", tc.size, got, tc.wantChunks) + } + }) + } +} + +// TestWsUploadInitBodyJSON checks that wsUploadInitBody serializes and +// deserializes correctly, including the omitempty MD5 field. +func TestWsUploadInitBodyJSON(t *testing.T) { + b := wsUploadInitBody{ + Type: "image", + Filename: "photo.jpg", + TotalSize: 1024, + TotalChunks: 1, + MD5: "deadbeefdeadbeef", + } + data, err := json.Marshal(b) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var got wsUploadInitBody + if err := json.Unmarshal(data, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got != b { + t.Errorf("round-trip mismatch: got %+v, want %+v", got, b) + } + + // MD5 must be omitted when empty. + b.MD5 = "" + data, _ = json.Marshal(b) + if contains(string(data), `"md5"`) { + t.Errorf("expected md5 key to be absent when empty, got: %s", data) + } +} + +// TestWsUploadChunkBodyBase64 checks that wsUploadChunkBody correctly +// round-trips arbitrary binary data through base64. +func TestWsUploadChunkBodyBase64(t *testing.T) { + raw := []byte("binary\x00data\xff\xfe") + body := wsUploadChunkBody{ + UploadID: "uid-test-1", + ChunkIndex: 0, + Base64Data: base64.StdEncoding.EncodeToString(raw), + } + encoded, err := json.Marshal(body) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var got wsUploadChunkBody + if err = json.Unmarshal(encoded, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + decoded, err := base64.StdEncoding.DecodeString(got.Base64Data) + if err != nil { + t.Fatalf("decode base64: %v", err) + } + if string(decoded) != string(raw) { + t.Errorf("base64 round-trip mismatch: got %v, want %v", decoded, raw) + } +} + +// TestWsUploadFinishResponseJSON checks that the finish response unmarshals +// the fields WeCom returns. +func TestWsUploadFinishResponseJSON(t *testing.T) { + raw := `{"type":"file","media_id":"MEDIAID_ABCDE","created_at":"1680000000"}` + var resp wsUploadFinishResponse + if err := json.Unmarshal([]byte(raw), &resp); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if resp.MediaID != "MEDIAID_ABCDE" { + t.Errorf("media_id: got %q, want %q", resp.MediaID, "MEDIAID_ABCDE") + } + if resp.Type != "file" { + t.Errorf("type: got %q, want %q", resp.Type, "file") + } + if resp.CreatedAt != "1680000000" { + t.Errorf("created_at: got %q, want %q", resp.CreatedAt, "1680000000") + } +} + +// contains is a helper to avoid importing "strings" for a single check. +func contains(s, sub string) bool { + return len(s) >= len(sub) && (s == sub || len(s) > 0 && containsHelper(s, sub)) +} + +func containsHelper(s, sub string) bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +}