初始化仓库
This commit is contained in:
@@ -0,0 +1,679 @@
|
||||
package weibo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultTokenEndpoint = "https://open-im.api.weibo.com/open/auth/ws_token"
|
||||
defaultWSEndpoint = "ws://open-im.api.weibo.com/ws/stream"
|
||||
|
||||
pingInterval = 30 * time.Second
|
||||
pongTimeout = 40 * time.Second
|
||||
reconnectDelay = 3 * time.Second
|
||||
maxReconnect = 10 * time.Second
|
||||
tokenRenewBuf = 60 * time.Second
|
||||
maxTextPerChunk = 2000
|
||||
maxSeenMessages = 1000
|
||||
)
|
||||
|
||||
func init() {
|
||||
core.RegisterPlatform("weibo", New)
|
||||
}
|
||||
|
||||
type replyContext struct {
|
||||
fromUserID string
|
||||
sessionKey string
|
||||
}
|
||||
|
||||
type Platform struct {
|
||||
name string
|
||||
appID string
|
||||
appSecret string
|
||||
tokenEndpoint string
|
||||
wsEndpoint string
|
||||
allowFrom string
|
||||
|
||||
handler core.MessageHandler
|
||||
|
||||
ws *websocket.Conn
|
||||
wsMu sync.Mutex
|
||||
connMu sync.Mutex
|
||||
|
||||
token string
|
||||
tokenExpiry time.Time
|
||||
tokenMu sync.Mutex
|
||||
uid string
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
seen map[string]struct{}
|
||||
seenMu sync.Mutex
|
||||
}
|
||||
|
||||
func New(opts map[string]any) (core.Platform, error) {
|
||||
name, _ := opts["name"].(string)
|
||||
if name == "" {
|
||||
name = "weibo"
|
||||
}
|
||||
appID, _ := opts["app_id"].(string)
|
||||
appSecret, _ := opts["app_secret"].(string)
|
||||
if appID == "" || appSecret == "" {
|
||||
return nil, fmt.Errorf("weibo: app_id and app_secret are required")
|
||||
}
|
||||
|
||||
tokenEndpoint, _ := opts["token_endpoint"].(string)
|
||||
if tokenEndpoint == "" {
|
||||
tokenEndpoint = defaultTokenEndpoint
|
||||
}
|
||||
wsEndpoint, _ := opts["ws_endpoint"].(string)
|
||||
if wsEndpoint == "" {
|
||||
wsEndpoint = defaultWSEndpoint
|
||||
}
|
||||
|
||||
allowFrom, _ := opts["allow_from"].(string)
|
||||
core.CheckAllowFrom(name, allowFrom)
|
||||
|
||||
return &Platform{
|
||||
name: name,
|
||||
appID: appID,
|
||||
appSecret: appSecret,
|
||||
tokenEndpoint: tokenEndpoint,
|
||||
wsEndpoint: wsEndpoint,
|
||||
allowFrom: allowFrom,
|
||||
seen: make(map[string]struct{}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (p *Platform) Name() string { return p.name }
|
||||
|
||||
func (p *Platform) Start(handler core.MessageHandler) error {
|
||||
p.handler = handler
|
||||
p.ctx, p.cancel = context.WithCancel(context.Background())
|
||||
|
||||
if _, err := p.refreshToken(); err != nil {
|
||||
return fmt.Errorf("weibo: initial token fetch: %w", err)
|
||||
}
|
||||
slog.Info(p.tag()+": authenticated", "uid", p.uid)
|
||||
|
||||
go p.connectLoop()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *Platform) Reply(ctx context.Context, rctx any, content string) error {
|
||||
return p.sendMessage(rctx, content)
|
||||
}
|
||||
|
||||
func (p *Platform) Send(ctx context.Context, rctx any, content string) error {
|
||||
return p.sendMessage(rctx, content)
|
||||
}
|
||||
|
||||
func (p *Platform) Stop() error {
|
||||
if p.cancel != nil {
|
||||
p.cancel()
|
||||
}
|
||||
p.wsMu.Lock()
|
||||
defer p.wsMu.Unlock()
|
||||
if p.ws != nil {
|
||||
p.ws.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Token management ---
|
||||
|
||||
type tokenResponse struct {
|
||||
Data struct {
|
||||
Token string `json:"token"`
|
||||
ExpireIn int64 `json:"expire_in"` // seconds
|
||||
UID json.RawMessage `json:"uid"`
|
||||
} `json:"data"`
|
||||
Error string `json:"error"`
|
||||
ErrorCode int `json:"error_code"`
|
||||
}
|
||||
|
||||
func (p *Platform) refreshToken() (string, error) {
|
||||
p.tokenMu.Lock()
|
||||
defer p.tokenMu.Unlock()
|
||||
|
||||
if p.token != "" && time.Now().Before(p.tokenExpiry.Add(-tokenRenewBuf)) {
|
||||
return p.token, nil
|
||||
}
|
||||
|
||||
body := fmt.Sprintf(`{"app_id":"%s","app_secret":"%s"}`, p.appID, p.appSecret)
|
||||
req, err := http.NewRequest("POST", p.tokenEndpoint, strings.NewReader(body))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("weibo: build token request: %w", err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("weibo: token request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("weibo: read token response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("weibo: token HTTP %d: %s", resp.StatusCode, string(raw))
|
||||
}
|
||||
|
||||
var tr tokenResponse
|
||||
if err := json.Unmarshal(raw, &tr); err != nil {
|
||||
return "", fmt.Errorf("weibo: parse token response: %w", err)
|
||||
}
|
||||
if tr.Data.Token == "" {
|
||||
return "", fmt.Errorf("weibo: empty token in response: %s", string(raw))
|
||||
}
|
||||
|
||||
p.token = tr.Data.Token
|
||||
p.tokenExpiry = time.Now().Add(time.Duration(tr.Data.ExpireIn) * time.Second)
|
||||
if len(tr.Data.UID) > 0 {
|
||||
p.uid = strings.Trim(string(tr.Data.UID), `"`)
|
||||
}
|
||||
|
||||
slog.Debug(p.tag()+": token refreshed", "expires_in", tr.Data.ExpireIn)
|
||||
return p.token, nil
|
||||
}
|
||||
|
||||
func (p *Platform) getToken() (string, error) {
|
||||
p.tokenMu.Lock()
|
||||
tok := p.token
|
||||
exp := p.tokenExpiry
|
||||
p.tokenMu.Unlock()
|
||||
|
||||
if tok != "" && time.Now().Before(exp.Add(-tokenRenewBuf)) {
|
||||
return tok, nil
|
||||
}
|
||||
return p.refreshToken()
|
||||
}
|
||||
|
||||
func (p *Platform) invalidateToken() {
|
||||
p.tokenMu.Lock()
|
||||
defer p.tokenMu.Unlock()
|
||||
p.token = ""
|
||||
p.tokenExpiry = time.Time{}
|
||||
}
|
||||
|
||||
// --- WebSocket connection ---
|
||||
|
||||
func (p *Platform) connectLoop() {
|
||||
delay := reconnectDelay
|
||||
for {
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
if err := p.connect(); err != nil {
|
||||
slog.Error(p.tag()+": connect failed", "error", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
return
|
||||
case <-time.After(delay):
|
||||
}
|
||||
|
||||
delay = min(delay*2, maxReconnect)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Platform) connect() error {
|
||||
p.connMu.Lock()
|
||||
defer p.connMu.Unlock()
|
||||
|
||||
tok, err := p.getToken()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
u, err := url.Parse(p.wsEndpoint)
|
||||
if err != nil {
|
||||
return fmt.Errorf("weibo: parse ws endpoint: %w", err)
|
||||
}
|
||||
q := u.Query()
|
||||
q.Set("app_id", p.appID)
|
||||
q.Set("token", tok)
|
||||
u.RawQuery = q.Encode()
|
||||
|
||||
ws, _, err := websocket.DefaultDialer.DialContext(p.ctx, u.String(), nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("weibo: ws dial: %w", err)
|
||||
}
|
||||
|
||||
p.wsMu.Lock()
|
||||
if p.ws != nil {
|
||||
p.ws.Close()
|
||||
}
|
||||
p.ws = ws
|
||||
p.wsMu.Unlock()
|
||||
|
||||
slog.Info(p.tag() + ": websocket connected")
|
||||
|
||||
go p.pingLoop(ws)
|
||||
p.readLoop(ws)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *Platform) pingLoop(ws *websocket.Conn) {
|
||||
ticker := time.NewTicker(pingInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-p.ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
p.wsMu.Lock()
|
||||
if p.ws != ws {
|
||||
p.wsMu.Unlock()
|
||||
return
|
||||
}
|
||||
err := ws.WriteJSON(map[string]string{"type": "ping"})
|
||||
p.wsMu.Unlock()
|
||||
if err != nil {
|
||||
slog.Debug(p.tag()+": ping write error", "error", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type wsMessage struct {
|
||||
Type string `json:"type"`
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
}
|
||||
|
||||
type messagePayload struct {
|
||||
MessageID string `json:"messageId"`
|
||||
FromUserID string `json:"fromUserId"`
|
||||
Text string `json:"text"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
Input []messageInputItem `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
type messageInputItem struct {
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []contentPart `json:"content"`
|
||||
}
|
||||
|
||||
type contentPart struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
Source *inputSource `json:"source,omitempty"`
|
||||
FileName string `json:"filename,omitempty"`
|
||||
}
|
||||
|
||||
type inputSource struct {
|
||||
Type string `json:"type"`
|
||||
MediaType string `json:"media_type"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
var supportedImageMIME = map[string]bool{
|
||||
"image/jpeg": true,
|
||||
"image/png": true,
|
||||
"image/gif": true,
|
||||
"image/webp": true,
|
||||
}
|
||||
|
||||
const (
|
||||
maxInboundImageBytes = 10 * 1024 * 1024
|
||||
maxInboundFileBytes = 5 * 1024 * 1024
|
||||
)
|
||||
|
||||
func (p *Platform) readLoop(ws *websocket.Conn) {
|
||||
ws.SetPongHandler(func(string) error {
|
||||
ws.SetReadDeadline(time.Now().Add(pongTimeout))
|
||||
return nil
|
||||
})
|
||||
ws.SetReadDeadline(time.Now().Add(pongTimeout))
|
||||
|
||||
for {
|
||||
_, raw, err := ws.ReadMessage()
|
||||
if err != nil {
|
||||
if p.ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if websocket.IsCloseError(err, 4002) {
|
||||
slog.Warn(p.tag() + ": invalid token, clearing cache")
|
||||
p.invalidateToken()
|
||||
}
|
||||
slog.Debug(p.tag()+": ws read error", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
ws.SetReadDeadline(time.Now().Add(pongTimeout))
|
||||
|
||||
var msg wsMessage
|
||||
if err := json.Unmarshal(raw, &msg); err != nil {
|
||||
slog.Debug(p.tag()+": ws unmarshal error", "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case "pong":
|
||||
// heartbeat response, already handled via deadline reset
|
||||
case "message":
|
||||
if msg.Payload != nil {
|
||||
p.handleInbound(msg.Payload)
|
||||
}
|
||||
default:
|
||||
slog.Debug(p.tag()+": unhandled ws message type", "type", msg.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Platform) handleInbound(raw json.RawMessage) {
|
||||
var payload messagePayload
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
slog.Error(p.tag()+": parse inbound payload", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
text, images, files := normalizeInboundInput(payload)
|
||||
hasText := strings.TrimSpace(text) != ""
|
||||
hasAttachments := len(images) > 0 || len(files) > 0
|
||||
|
||||
if payload.FromUserID == "" || (!hasText && !hasAttachments) {
|
||||
return
|
||||
}
|
||||
|
||||
msgID := payload.MessageID
|
||||
if msgID == "" {
|
||||
msgID = fmt.Sprintf("%x", sha256.Sum256(raw))[:16]
|
||||
}
|
||||
|
||||
if p.isDuplicate(msgID) {
|
||||
return
|
||||
}
|
||||
|
||||
userID := payload.FromUserID
|
||||
if !core.AllowList(p.allowFrom, userID) {
|
||||
slog.Debug(p.tag()+": message from unauthorized user", "user", userID)
|
||||
return
|
||||
}
|
||||
|
||||
sessionKey := p.name + ":" + userID + ":" + userID
|
||||
rctx := replyContext{
|
||||
fromUserID: userID,
|
||||
sessionKey: sessionKey,
|
||||
}
|
||||
|
||||
msg := &core.Message{
|
||||
SessionKey: sessionKey,
|
||||
Platform: p.name,
|
||||
MessageID: msgID,
|
||||
UserID: userID,
|
||||
UserName: userID,
|
||||
Content: text,
|
||||
Images: images,
|
||||
Files: files,
|
||||
ReplyCtx: rctx,
|
||||
}
|
||||
|
||||
if hasAttachments {
|
||||
slog.Debug(p.tag()+": inbound with attachments",
|
||||
"user", userID, "images", len(images), "files", len(files))
|
||||
}
|
||||
|
||||
p.handler(p, msg)
|
||||
}
|
||||
|
||||
func normalizeInboundInput(payload messagePayload) (string, []core.ImageAttachment, []core.FileAttachment) {
|
||||
var textParts []string
|
||||
var images []core.ImageAttachment
|
||||
var files []core.FileAttachment
|
||||
|
||||
for _, item := range payload.Input {
|
||||
if item.Type != "message" || item.Role != "user" {
|
||||
continue
|
||||
}
|
||||
for _, part := range item.Content {
|
||||
switch part.Type {
|
||||
case "input_text":
|
||||
if part.Text != "" {
|
||||
textParts = append(textParts, part.Text)
|
||||
}
|
||||
case "input_image":
|
||||
if part.Source == nil || part.Source.Data == "" {
|
||||
continue
|
||||
}
|
||||
if !supportedImageMIME[part.Source.MediaType] {
|
||||
slog.Warn("weibo: unsupported inbound image mime", "mime", part.Source.MediaType)
|
||||
continue
|
||||
}
|
||||
data, err := base64.StdEncoding.DecodeString(part.Source.Data)
|
||||
if err != nil {
|
||||
slog.Warn("weibo: decode inbound image base64", "error", err)
|
||||
continue
|
||||
}
|
||||
if len(data) == 0 || len(data) > maxInboundImageBytes {
|
||||
continue
|
||||
}
|
||||
images = append(images, core.ImageAttachment{
|
||||
MimeType: part.Source.MediaType,
|
||||
Data: data,
|
||||
FileName: part.FileName,
|
||||
})
|
||||
case "input_file":
|
||||
if part.Source == nil || part.Source.Data == "" {
|
||||
continue
|
||||
}
|
||||
data, err := base64.StdEncoding.DecodeString(part.Source.Data)
|
||||
if err != nil {
|
||||
slog.Warn("weibo: decode inbound file base64", "error", err)
|
||||
continue
|
||||
}
|
||||
if len(data) == 0 || len(data) > maxInboundFileBytes {
|
||||
continue
|
||||
}
|
||||
files = append(files, core.FileAttachment{
|
||||
MimeType: part.Source.MediaType,
|
||||
Data: data,
|
||||
FileName: part.FileName,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
text := payload.Text
|
||||
if len(textParts) > 0 {
|
||||
text = strings.Join(textParts, "\n")
|
||||
}
|
||||
|
||||
return text, images, files
|
||||
}
|
||||
|
||||
// --- Sending ---
|
||||
|
||||
type sendPayload struct {
|
||||
ToUserID string `json:"toUserId"`
|
||||
Text string `json:"text"`
|
||||
MessageID string `json:"messageId"`
|
||||
ChunkID int `json:"chunkId"`
|
||||
Done bool `json:"done"`
|
||||
Input []messageInputItem `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
func (p *Platform) sendMessage(rctx any, content string) error {
|
||||
rc, ok := rctx.(replyContext)
|
||||
if !ok {
|
||||
return fmt.Errorf("weibo: invalid reply context type: %T", rctx)
|
||||
}
|
||||
|
||||
chunks := splitText(content, maxTextPerChunk)
|
||||
msgID := fmt.Sprintf("out-%s-%d", rc.fromUserID, time.Now().UnixMilli())
|
||||
|
||||
for i, chunk := range chunks {
|
||||
env := map[string]any{
|
||||
"type": "send_message",
|
||||
"payload": sendPayload{
|
||||
ToUserID: rc.fromUserID,
|
||||
Text: chunk,
|
||||
MessageID: msgID,
|
||||
ChunkID: i,
|
||||
Done: i == len(chunks)-1,
|
||||
},
|
||||
}
|
||||
if err := p.writeWS(env); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *Platform) SendImage(_ context.Context, rctx any, img core.ImageAttachment) error {
|
||||
rc, ok := rctx.(replyContext)
|
||||
if !ok {
|
||||
return fmt.Errorf("weibo: invalid reply context type: %T", rctx)
|
||||
}
|
||||
|
||||
b64 := base64.StdEncoding.EncodeToString(img.Data)
|
||||
mime := img.MimeType
|
||||
if mime == "" {
|
||||
mime = "image/png"
|
||||
}
|
||||
fname := img.FileName
|
||||
if fname == "" {
|
||||
fname = "image"
|
||||
}
|
||||
|
||||
msgID := fmt.Sprintf("img-%s-%d", rc.fromUserID, time.Now().UnixMilli())
|
||||
env := map[string]any{
|
||||
"type": "send_message",
|
||||
"payload": sendPayload{
|
||||
ToUserID: rc.fromUserID,
|
||||
MessageID: msgID,
|
||||
Done: true,
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []contentPart{{
|
||||
Type: "input_image",
|
||||
FileName: fname,
|
||||
Source: &inputSource{Type: "base64", MediaType: mime, Data: b64},
|
||||
}},
|
||||
}},
|
||||
},
|
||||
}
|
||||
slog.Debug(p.tag()+": sending image", "to", rc.fromUserID, "name", fname, "size", len(img.Data))
|
||||
return p.writeWS(env)
|
||||
}
|
||||
|
||||
func (p *Platform) SendFile(_ context.Context, rctx any, file core.FileAttachment) error {
|
||||
rc, ok := rctx.(replyContext)
|
||||
if !ok {
|
||||
return fmt.Errorf("weibo: invalid reply context type: %T", rctx)
|
||||
}
|
||||
|
||||
b64 := base64.StdEncoding.EncodeToString(file.Data)
|
||||
mime := file.MimeType
|
||||
if mime == "" {
|
||||
mime = "application/octet-stream"
|
||||
}
|
||||
fname := file.FileName
|
||||
if fname == "" {
|
||||
fname = "attachment"
|
||||
}
|
||||
|
||||
msgID := fmt.Sprintf("file-%s-%d", rc.fromUserID, time.Now().UnixMilli())
|
||||
env := map[string]any{
|
||||
"type": "send_message",
|
||||
"payload": sendPayload{
|
||||
ToUserID: rc.fromUserID,
|
||||
MessageID: msgID,
|
||||
Done: true,
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []contentPart{{
|
||||
Type: "input_file",
|
||||
FileName: fname,
|
||||
Source: &inputSource{Type: "base64", MediaType: mime, Data: b64},
|
||||
}},
|
||||
}},
|
||||
},
|
||||
}
|
||||
slog.Debug(p.tag()+": sending file", "to", rc.fromUserID, "name", fname, "size", len(file.Data))
|
||||
return p.writeWS(env)
|
||||
}
|
||||
|
||||
func (p *Platform) writeWS(data any) error {
|
||||
// gorilla/websocket only allows one concurrent writer; wsMu must guard the
|
||||
// full WriteJSON call (pingLoop already follows this pattern), otherwise
|
||||
// concurrent sendMessage / SendImage / SendFile calls interleave frames
|
||||
// on the wire.
|
||||
p.wsMu.Lock()
|
||||
defer p.wsMu.Unlock()
|
||||
if p.ws == nil {
|
||||
return fmt.Errorf("weibo: not connected")
|
||||
}
|
||||
if err := p.ws.WriteJSON(data); err != nil {
|
||||
return fmt.Errorf("weibo: ws send: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
func (p *Platform) tag() string { return p.name }
|
||||
|
||||
func (p *Platform) isDuplicate(msgID string) bool {
|
||||
p.seenMu.Lock()
|
||||
defer p.seenMu.Unlock()
|
||||
if _, ok := p.seen[msgID]; ok {
|
||||
return true
|
||||
}
|
||||
if len(p.seen) >= maxSeenMessages {
|
||||
// prune half
|
||||
i := 0
|
||||
for k := range p.seen {
|
||||
if i >= maxSeenMessages/2 {
|
||||
break
|
||||
}
|
||||
delete(p.seen, k)
|
||||
i++
|
||||
}
|
||||
}
|
||||
p.seen[msgID] = struct{}{}
|
||||
return false
|
||||
}
|
||||
|
||||
func splitText(text string, limit int) []string {
|
||||
runes := []rune(text)
|
||||
if len(runes) <= limit {
|
||||
return []string{text}
|
||||
}
|
||||
var chunks []string
|
||||
for len(runes) > 0 {
|
||||
end := limit
|
||||
if end > len(runes) {
|
||||
end = len(runes)
|
||||
}
|
||||
chunks = append(chunks, string(runes[:end]))
|
||||
runes = runes[end:]
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
@@ -0,0 +1,818 @@
|
||||
package weibo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func TestNew_RequiredFields(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts map[string]any
|
||||
wantErr bool
|
||||
}{
|
||||
{"missing both", map[string]any{}, true},
|
||||
{"missing app_secret", map[string]any{"app_id": "id"}, true},
|
||||
{"missing app_id", map[string]any{"app_secret": "secret"}, true},
|
||||
{"valid", map[string]any{"app_id": "id", "app_secret": "secret"}, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
p, err := New(tt.opts)
|
||||
if tt.wantErr {
|
||||
if err == nil {
|
||||
t.Error("expected error")
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if p.Name() != "weibo" {
|
||||
t.Errorf("name = %q, want %q", p.Name(), "weibo")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNew_CustomName(t *testing.T) {
|
||||
p, err := New(map[string]any{
|
||||
"app_id": "id",
|
||||
"app_secret": "secret",
|
||||
"name": "my-weibo",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if p.Name() != "my-weibo" {
|
||||
t.Errorf("name = %q, want %q", p.Name(), "my-weibo")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNew_CustomEndpoints(t *testing.T) {
|
||||
p, err := New(map[string]any{
|
||||
"app_id": "id",
|
||||
"app_secret": "secret",
|
||||
"token_endpoint": "https://custom.example.com/token",
|
||||
"ws_endpoint": "ws://custom.example.com/ws",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
plat := p.(*Platform)
|
||||
if plat.tokenEndpoint != "https://custom.example.com/token" {
|
||||
t.Errorf("tokenEndpoint = %q", plat.tokenEndpoint)
|
||||
}
|
||||
if plat.wsEndpoint != "ws://custom.example.com/ws" {
|
||||
t.Errorf("wsEndpoint = %q", plat.wsEndpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
limit int
|
||||
chunks int
|
||||
}{
|
||||
{"short", "hello", 100, 1},
|
||||
{"exact", "abcde", 5, 1},
|
||||
{"split", "abcdefgh", 3, 3},
|
||||
{"empty", "", 10, 1},
|
||||
{"unicode", "你好世界测试", 3, 2},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := splitText(tt.text, tt.limit)
|
||||
if len(result) != tt.chunks {
|
||||
t.Errorf("splitText(%q, %d) = %d chunks, want %d", tt.text, tt.limit, len(result), tt.chunks)
|
||||
}
|
||||
joined := strings.Join(result, "")
|
||||
if joined != tt.text {
|
||||
t.Errorf("joined = %q, want %q", joined, tt.text)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDuplicate(t *testing.T) {
|
||||
p := &Platform{seen: make(map[string]struct{})}
|
||||
|
||||
if p.isDuplicate("msg1") {
|
||||
t.Error("first occurrence should not be duplicate")
|
||||
}
|
||||
if !p.isDuplicate("msg1") {
|
||||
t.Error("second occurrence should be duplicate")
|
||||
}
|
||||
if p.isDuplicate("msg2") {
|
||||
t.Error("different message should not be duplicate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDuplicate_Prune(t *testing.T) {
|
||||
p := &Platform{seen: make(map[string]struct{})}
|
||||
|
||||
for i := 0; i < maxSeenMessages+100; i++ {
|
||||
p.isDuplicate(strings.Repeat("x", 10) + string(rune(i)))
|
||||
}
|
||||
if len(p.seen) > maxSeenMessages {
|
||||
t.Errorf("seen map should be pruned, got %d entries", len(p.seen))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
var mu sync.Mutex
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
mu.Lock()
|
||||
received = msg
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "test-123",
|
||||
FromUserID: "user1",
|
||||
Text: "hello world",
|
||||
Timestamp: 1234567890,
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if received == nil {
|
||||
t.Fatal("handler not called")
|
||||
}
|
||||
if received.SessionKey != "weibo:user1:user1" {
|
||||
t.Errorf("sessionKey = %q", received.SessionKey)
|
||||
}
|
||||
if received.Content != "hello world" {
|
||||
t.Errorf("content = %q", received.Content)
|
||||
}
|
||||
if received.UserID != "user1" {
|
||||
t.Errorf("userID = %q", received.UserID)
|
||||
}
|
||||
if received.MessageID != "test-123" {
|
||||
t.Errorf("messageID = %q", received.MessageID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_AllowList(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "user2,user3",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
called := false
|
||||
p.handler = func(_ core.Platform, _ *core.Message) {
|
||||
called = true
|
||||
}
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "blocked-1",
|
||||
FromUserID: "user1",
|
||||
Text: "hello",
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
if called {
|
||||
t.Error("handler should not be called for unauthorized user")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_EmptyText(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
called := false
|
||||
p.handler = func(_ core.Platform, _ *core.Message) {
|
||||
called = true
|
||||
}
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "empty-1",
|
||||
FromUserID: "user1",
|
||||
Text: "",
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
if called {
|
||||
t.Error("handler should not be called for empty text without attachments")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_WithImage(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
var mu sync.Mutex
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
mu.Lock()
|
||||
received = msg
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
imgData := []byte("fake-png-data")
|
||||
b64 := base64.StdEncoding.EncodeToString(imgData)
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "img-1",
|
||||
FromUserID: "user1",
|
||||
Text: "check this image",
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "user",
|
||||
Content: []contentPart{{
|
||||
Type: "input_image",
|
||||
FileName: "photo.png",
|
||||
Source: &inputSource{Type: "base64", MediaType: "image/png", Data: b64},
|
||||
}},
|
||||
}},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if received == nil {
|
||||
t.Fatal("handler not called")
|
||||
}
|
||||
if received.Content != "check this image" {
|
||||
t.Errorf("content = %q", received.Content)
|
||||
}
|
||||
if len(received.Images) != 1 {
|
||||
t.Fatalf("images = %d, want 1", len(received.Images))
|
||||
}
|
||||
if received.Images[0].MimeType != "image/png" {
|
||||
t.Errorf("image mime = %q", received.Images[0].MimeType)
|
||||
}
|
||||
if received.Images[0].FileName != "photo.png" {
|
||||
t.Errorf("image filename = %q", received.Images[0].FileName)
|
||||
}
|
||||
if string(received.Images[0].Data) != "fake-png-data" {
|
||||
t.Errorf("image data mismatch")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_WithFile(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
var mu sync.Mutex
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
mu.Lock()
|
||||
received = msg
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
fileData := []byte("hello world pdf content")
|
||||
b64 := base64.StdEncoding.EncodeToString(fileData)
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "file-1",
|
||||
FromUserID: "user1",
|
||||
Text: "",
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "user",
|
||||
Content: []contentPart{
|
||||
{Type: "input_text", Text: "here is my file"},
|
||||
{
|
||||
Type: "input_file",
|
||||
FileName: "doc.pdf",
|
||||
Source: &inputSource{Type: "base64", MediaType: "application/pdf", Data: b64},
|
||||
},
|
||||
},
|
||||
}},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if received == nil {
|
||||
t.Fatal("handler not called")
|
||||
}
|
||||
if received.Content != "here is my file" {
|
||||
t.Errorf("content = %q, want %q", received.Content, "here is my file")
|
||||
}
|
||||
if len(received.Files) != 1 {
|
||||
t.Fatalf("files = %d, want 1", len(received.Files))
|
||||
}
|
||||
if received.Files[0].MimeType != "application/pdf" {
|
||||
t.Errorf("file mime = %q", received.Files[0].MimeType)
|
||||
}
|
||||
if received.Files[0].FileName != "doc.pdf" {
|
||||
t.Errorf("file name = %q", received.Files[0].FileName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_ImageOnlyNoText(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
received = msg
|
||||
}
|
||||
|
||||
imgData := []byte("image-bytes")
|
||||
b64 := base64.StdEncoding.EncodeToString(imgData)
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "imgonly-1",
|
||||
FromUserID: "user1",
|
||||
Text: "",
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "user",
|
||||
Content: []contentPart{{
|
||||
Type: "input_image",
|
||||
Source: &inputSource{Type: "base64", MediaType: "image/jpeg", Data: b64},
|
||||
}},
|
||||
}},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
if received == nil {
|
||||
t.Fatal("handler should be called for image-only message")
|
||||
}
|
||||
if len(received.Images) != 1 {
|
||||
t.Errorf("images = %d, want 1", len(received.Images))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_UnsupportedImageMime(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
received = msg
|
||||
}
|
||||
|
||||
b64 := base64.StdEncoding.EncodeToString([]byte("bmp-data"))
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "bmp-1",
|
||||
FromUserID: "user1",
|
||||
Text: "a bmp image",
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "user",
|
||||
Content: []contentPart{{
|
||||
Type: "input_image",
|
||||
Source: &inputSource{Type: "base64", MediaType: "image/bmp", Data: b64},
|
||||
}},
|
||||
}},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
if received == nil {
|
||||
t.Fatal("handler should be called for text content")
|
||||
}
|
||||
if len(received.Images) != 0 {
|
||||
t.Errorf("unsupported image should be filtered, got %d images", len(received.Images))
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleInbound_InputTextOverridesPayloadText(t *testing.T) {
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
allowFrom: "*",
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
var received *core.Message
|
||||
p.handler = func(_ core.Platform, msg *core.Message) {
|
||||
received = msg
|
||||
}
|
||||
|
||||
payload := messagePayload{
|
||||
MessageID: "override-1",
|
||||
FromUserID: "user1",
|
||||
Text: "payload text",
|
||||
Input: []messageInputItem{{
|
||||
Type: "message",
|
||||
Role: "user",
|
||||
Content: []contentPart{
|
||||
{Type: "input_text", Text: "input part 1"},
|
||||
{Type: "input_text", Text: "input part 2"},
|
||||
},
|
||||
}},
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
p.handleInbound(raw)
|
||||
|
||||
if received == nil {
|
||||
t.Fatal("handler not called")
|
||||
}
|
||||
if received.Content != "input part 1\ninput part 2" {
|
||||
t.Errorf("content = %q, want joined input_text", received.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeInboundInput_SkipsNonUserRole(t *testing.T) {
|
||||
payload := messagePayload{
|
||||
FromUserID: "user1",
|
||||
Text: "fallback",
|
||||
Input: []messageInputItem{
|
||||
{
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []contentPart{
|
||||
{Type: "input_text", Text: "should be ignored"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
text, images, files := normalizeInboundInput(payload)
|
||||
if text != "fallback" {
|
||||
t.Errorf("text = %q, want fallback", text)
|
||||
}
|
||||
if len(images) != 0 || len(files) != 0 {
|
||||
t.Error("should have no attachments from non-user role")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshToken(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "POST" {
|
||||
t.Errorf("method = %s, want POST", r.Method)
|
||||
}
|
||||
json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": map[string]any{
|
||||
"token": "test-token-abc",
|
||||
"expire_in": 3600,
|
||||
"uid": 12345,
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
p := &Platform{
|
||||
appID: "test-app",
|
||||
appSecret: "test-secret",
|
||||
tokenEndpoint: ts.URL,
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
tok, err := p.refreshToken()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tok != "test-token-abc" {
|
||||
t.Errorf("token = %q, want %q", tok, "test-token-abc")
|
||||
}
|
||||
if p.uid != "12345" {
|
||||
t.Errorf("uid = %q, want %q", p.uid, "12345")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendMessage(t *testing.T) {
|
||||
upgrader := websocket.Upgrader{}
|
||||
gotMsg := make(chan map[string]any, 1)
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
c, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
for {
|
||||
_, msg, err := c.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var m map[string]any
|
||||
json.Unmarshal(msg, &m)
|
||||
gotMsg <- m
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(ts.URL, "http")
|
||||
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ws.Close()
|
||||
|
||||
p := &Platform{
|
||||
name: "weibo",
|
||||
ws: ws,
|
||||
seen: make(map[string]struct{}),
|
||||
}
|
||||
|
||||
rctx := replyContext{fromUserID: "user1", sessionKey: "weibo:user1:user1"}
|
||||
err = p.sendMessage(rctx, "short message")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
select {
|
||||
case m := <-gotMsg:
|
||||
if m["type"] != "send_message" {
|
||||
t.Errorf("type = %v", m["type"])
|
||||
}
|
||||
payload := m["payload"].(map[string]any)
|
||||
if payload["toUserId"] != "user1" {
|
||||
t.Errorf("toUserId = %v", payload["toUserId"])
|
||||
}
|
||||
if payload["text"] != "short message" {
|
||||
t.Errorf("text = %v", payload["text"])
|
||||
}
|
||||
if payload["done"] != true {
|
||||
t.Errorf("done = %v", payload["done"])
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out waiting for message")
|
||||
}
|
||||
}
|
||||
|
||||
func newWSTestPlatform(t *testing.T) (*Platform, chan map[string]any) {
|
||||
t.Helper()
|
||||
upgrader := websocket.Upgrader{}
|
||||
gotMsg := make(chan map[string]any, 5)
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
c, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
for {
|
||||
_, msg, err := c.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var m map[string]any
|
||||
json.Unmarshal(msg, &m)
|
||||
gotMsg <- m
|
||||
}
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(ts.URL, "http")
|
||||
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { ws.Close() })
|
||||
|
||||
p := &Platform{name: "weibo", ws: ws, seen: make(map[string]struct{})}
|
||||
return p, gotMsg
|
||||
}
|
||||
|
||||
func TestSendImage(t *testing.T) {
|
||||
p, gotMsg := newWSTestPlatform(t)
|
||||
|
||||
rctx := replyContext{fromUserID: "user1", sessionKey: "weibo:user1:user1"}
|
||||
imgData := []byte("fake-image-bytes")
|
||||
|
||||
err := p.SendImage(context.Background(), rctx, core.ImageAttachment{
|
||||
MimeType: "image/png",
|
||||
Data: imgData,
|
||||
FileName: "screenshot.png",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
select {
|
||||
case m := <-gotMsg:
|
||||
if m["type"] != "send_message" {
|
||||
t.Errorf("type = %v", m["type"])
|
||||
}
|
||||
payload := m["payload"].(map[string]any)
|
||||
if payload["toUserId"] != "user1" {
|
||||
t.Errorf("toUserId = %v", payload["toUserId"])
|
||||
}
|
||||
if payload["done"] != true {
|
||||
t.Errorf("done = %v", payload["done"])
|
||||
}
|
||||
|
||||
input, ok := payload["input"].([]any)
|
||||
if !ok || len(input) == 0 {
|
||||
t.Fatal("input missing or empty")
|
||||
}
|
||||
item := input[0].(map[string]any)
|
||||
if item["role"] != "assistant" {
|
||||
t.Errorf("role = %v", item["role"])
|
||||
}
|
||||
content := item["content"].([]any)
|
||||
part := content[0].(map[string]any)
|
||||
if part["type"] != "input_image" {
|
||||
t.Errorf("part type = %v", part["type"])
|
||||
}
|
||||
if part["filename"] != "screenshot.png" {
|
||||
t.Errorf("filename = %v", part["filename"])
|
||||
}
|
||||
src := part["source"].(map[string]any)
|
||||
if src["media_type"] != "image/png" {
|
||||
t.Errorf("media_type = %v", src["media_type"])
|
||||
}
|
||||
decoded, _ := base64.StdEncoding.DecodeString(src["data"].(string))
|
||||
if string(decoded) != string(imgData) {
|
||||
t.Error("image data mismatch after round-trip")
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendFile(t *testing.T) {
|
||||
p, gotMsg := newWSTestPlatform(t)
|
||||
|
||||
rctx := replyContext{fromUserID: "user1", sessionKey: "weibo:user1:user1"}
|
||||
fileData := []byte("pdf-content-here")
|
||||
|
||||
err := p.SendFile(context.Background(), rctx, core.FileAttachment{
|
||||
MimeType: "application/pdf",
|
||||
Data: fileData,
|
||||
FileName: "report.pdf",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
select {
|
||||
case m := <-gotMsg:
|
||||
payload := m["payload"].(map[string]any)
|
||||
input := payload["input"].([]any)
|
||||
item := input[0].(map[string]any)
|
||||
content := item["content"].([]any)
|
||||
part := content[0].(map[string]any)
|
||||
if part["type"] != "input_file" {
|
||||
t.Errorf("part type = %v", part["type"])
|
||||
}
|
||||
if part["filename"] != "report.pdf" {
|
||||
t.Errorf("filename = %v", part["filename"])
|
||||
}
|
||||
src := part["source"].(map[string]any)
|
||||
if src["media_type"] != "application/pdf" {
|
||||
t.Errorf("media_type = %v", src["media_type"])
|
||||
}
|
||||
decoded, _ := base64.StdEncoding.DecodeString(src["data"].(string))
|
||||
if string(decoded) != string(fileData) {
|
||||
t.Error("file data mismatch")
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("timed out")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendImage_NotConnected(t *testing.T) {
|
||||
p := &Platform{name: "weibo", seen: make(map[string]struct{})}
|
||||
rctx := replyContext{fromUserID: "u1"}
|
||||
|
||||
err := p.SendImage(context.Background(), rctx, core.ImageAttachment{Data: []byte("x")})
|
||||
if err == nil {
|
||||
t.Error("expected error when not connected")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "not connected") {
|
||||
t.Errorf("error = %q, want 'not connected'", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendFile_InvalidContext(t *testing.T) {
|
||||
p := &Platform{name: "weibo", seen: make(map[string]struct{})}
|
||||
err := p.SendFile(context.Background(), "invalid", core.FileAttachment{Data: []byte("x")})
|
||||
if err == nil {
|
||||
t.Error("expected error for invalid context")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInterfaceCompliance(t *testing.T) {
|
||||
var _ core.ImageSender = (*Platform)(nil)
|
||||
var _ core.FileSender = (*Platform)(nil)
|
||||
}
|
||||
|
||||
// TestWriteWS_ConcurrentSendsSerialized verifies that writeWS serializes
|
||||
// concurrent callers as gorilla/websocket requires (one writer at a time).
|
||||
// Without the wsMu fix, parallel WriteJSON calls race on the underlying
|
||||
// Conn.writer field (caught by go test -race) and may interleave frames.
|
||||
func TestWriteWS_ConcurrentSendsSerialized(t *testing.T) {
|
||||
upgrader := websocket.Upgrader{}
|
||||
const n = 50
|
||||
gotMsg := make(chan map[string]any, n*2)
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
c, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer c.Close()
|
||||
for {
|
||||
_, msg, err := c.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var m map[string]any
|
||||
if err := json.Unmarshal(msg, &m); err != nil {
|
||||
gotMsg <- map[string]any{"_parse_error": err.Error(), "_raw": string(msg)}
|
||||
continue
|
||||
}
|
||||
gotMsg <- m
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(ts.URL, "http")
|
||||
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer ws.Close()
|
||||
|
||||
p := &Platform{name: "weibo", ws: ws, seen: make(map[string]struct{})}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, n)
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
rctx := replyContext{fromUserID: fmt.Sprintf("u%d", i)}
|
||||
if err := p.sendMessage(rctx, fmt.Sprintf("m%d", i)); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
|
||||
for err := range errs {
|
||||
t.Errorf("concurrent sendMessage: %v", err)
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
deadline := time.After(3 * time.Second)
|
||||
for len(seen) < n {
|
||||
select {
|
||||
case m := <-gotMsg:
|
||||
if pe, ok := m["_parse_error"]; ok {
|
||||
t.Fatalf("server got malformed JSON frame (concurrent write interleaved): %v raw=%q", pe, m["_raw"])
|
||||
}
|
||||
payload, ok := m["payload"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("frame missing payload: %v", m)
|
||||
}
|
||||
to, _ := payload["toUserId"].(string)
|
||||
if to == "" {
|
||||
t.Fatalf("frame missing toUserId: %v", payload)
|
||||
}
|
||||
if seen[to] {
|
||||
t.Fatalf("duplicate frame for %s", to)
|
||||
}
|
||||
seen[to] = true
|
||||
case <-deadline:
|
||||
t.Fatalf("only got %d of %d messages within 3s", len(seen), n)
|
||||
}
|
||||
}
|
||||
if len(seen) != n {
|
||||
t.Fatalf("got %d unique frames, want %d", len(seen), n)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user