feat(upload): implement WeCom media upload via WebSocket API

This commit is contained in:
Zhang Rui 2026-03-19 22:04:13 +08:00
parent 9a25fad20a
commit e46ef82aac
6 changed files with 550 additions and 42 deletions

View file

@ -193,6 +193,7 @@ func registerSharedTools(
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
Content: content, Content: content,
Final: true,
}) })
}) })
agent.Tools.Register(messageTool) agent.Tools.Register(messageTool)
@ -315,6 +316,7 @@ func (al *AgentLoop) Run(ctx context.Context) error {
Channel: msg.Channel, Channel: msg.Channel,
ChatID: msg.ChatID, ChatID: msg.ChatID,
Content: response, Content: response,
Final: true,
}) })
logger.InfoCF("agent", "Published outbound response", logger.InfoCF("agent", "Published outbound response",
map[string]any{ map[string]any{

View file

@ -64,7 +64,7 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
return nil 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) logger.WarnCF("agent", "MCP is enabled but no servers are configured, skipping MCP initialization", nil)
return nil return nil
} }

View file

@ -34,6 +34,7 @@ type OutboundMessage struct {
ChatID string `json:"chat_id"` ChatID string `json:"chat_id"`
Content string `json:"content"` Content string `json:"content"`
ReplyToMessageID string `json:"reply_to_message_id,omitempty"` ReplyToMessageID string `json:"reply_to_message_id,omitempty"`
Final bool `json:"final,omitempty"`
} }
// MediaPart describes a single media attachment to send. // MediaPart describes a single media attachment to send.

View file

@ -38,9 +38,7 @@ const (
wsInitialReconnect = time.Second wsInitialReconnect = time.Second
// WeCom requires finish=true within 6 minutes of the first stream frame. // 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. // 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 caps the time we spend downloading an inbound image.
@ -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. // and there will be no matching entry in reqStates; fall through to proactive push.
task, route, ok := c.getReqState(msg.ChatID) task, route, ok := c.getReqState(msg.ChatID)
if !ok { 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. // No req_id record found — this is a cron/scheduler-originated message.
// Send it as a proactive markdown push using the chat ID directly. // 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)", 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 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) { if time.Now().Before(route.ReadyAt) {
// Keep using aibot_respond_msg within stream window; do not proactively // Keep using aibot_respond_msg within stream window; do not proactively
// push unless wsStreamMaxDuration has elapsed. // push unless wsStreamMaxDuration has elapsed.
@ -318,6 +327,19 @@ func (c *WeComAIBotWSChannel) Send(ctx context.Context, msg bus.OutboundMessage)
return nil 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 // 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 // against task.ctx.Done() (which fires when the task is canceled by a new
// incoming message, but the response must still be sent). // incoming message, but the response must still be sent).
@ -852,37 +874,34 @@ func (c *WeComAIBotWSChannel) dispatchWSAgentTask(
c.HandleMessage(taskCtx, peer, reqID, userID, reqID, c.HandleMessage(taskCtx, peer, reqID, userID, reqID,
content, mediaRefs, metadata, sender) content, mediaRefs, metadata, sender)
// Wait for the agent reply. While waiting, send periodic finish=false // Wait for the agent reply.
// hints so the user knows processing is still in progress.
// WeCom requires finish=true within 6 minutes of the first stream frame; // WeCom requires finish=true within 6 minutes of the first stream frame;
// wsStreamMaxDuration enforces that limit with a safety margin. // 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) deadlineTimer := time.NewTimer(wsStreamMaxDuration)
defer deadlineTimer.Stop() defer deadlineTimer.Stop()
tickCount := 0
for { for {
select { select {
case answer := <-task.answerCh: 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) chunks := splitWSContent(answer, wsStreamMaxContentBytes)
for i, chunk := range chunks { 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) c.deleteReqState(reqID)
return 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: case <-deadlineTimer.C:
logger.WarnCF("wecom_aibot", logger.WarnCF("wecom_aibot",
"Stream response deadline reached, closing stream; late reply will be pushed", "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. // 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 { func (c *WeComAIBotWSChannel) writeWSAndWait(cmd wsCommand, timeout time.Duration) error {
if cmd.Headers.ReqID == "" { _, err := c.callWSCommand(cmd, timeout)
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 return err
}
if resp.ErrCode != 0 {
return fmt.Errorf("%s rejected (errcode=%d): %s", cmd.Cmd, resp.ErrCode, resp.ErrMsg)
}
return nil
} }
// cancelAllTasks cancels every pending agent task; called when the connection drops. // cancelAllTasks cancels every pending agent task; called when the connection drops.

View file

@ -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
}

View file

@ -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
}