Files
cc-connect/core/streaming_test.go
T
chenhg5andClaude Opus 4.6 e5d328d9d6 feat: add streaming preview and rate limiting support
Streaming Preview:
- Real-time message preview during agent response (like "typing" indicator)
- Supported on Telegram, Discord, and Feishu
- Configurable interval, min delta chars, and max preview length

Rate Limiting:
- Per-session sliding-window rate limiter
- Prevents users from flooding the bot with messages
- Configurable max messages and window size

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-06 07:02:37 +08:00

185 lines
4.4 KiB
Go

package core
import (
"context"
"sync"
"testing"
"time"
)
// mockUpdaterPlatform implements Platform + MessageUpdater + PreviewStarter.
type mockUpdaterPlatform struct {
stubPlatformEngine
mu sync.Mutex
messages []string // track all sent/updated messages
lastMsg string
}
func (m *mockUpdaterPlatform) SendPreviewStart(_ context.Context, _ any, content string) (any, error) {
m.mu.Lock()
defer m.mu.Unlock()
m.messages = append(m.messages, "start:"+content)
m.lastMsg = content
return "preview-handle", nil
}
func (m *mockUpdaterPlatform) UpdateMessage(_ context.Context, _ any, content string) error {
m.mu.Lock()
defer m.mu.Unlock()
m.messages = append(m.messages, "update:"+content)
m.lastMsg = content
return nil
}
func (m *mockUpdaterPlatform) getMessages() []string {
m.mu.Lock()
defer m.mu.Unlock()
out := make([]string, len(m.messages))
copy(out, m.messages)
return out
}
func TestStreamPreview_BasicFlow(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 100,
MinDeltaChars: 5,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
if !sp.canPreview() {
t.Fatal("should be able to preview")
}
sp.appendText("Hello ")
time.Sleep(150 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) == 0 {
t.Fatal("expected at least one message sent")
}
if msgs[0] != "start:Hello ▍" {
t.Errorf("first message = %q, want 'start:Hello ▍'", msgs[0])
}
}
func TestStreamPreview_ThrottlesUpdates(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 200,
MinDeltaChars: 5,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
// Rapid-fire small appends
for i := 0; i < 10; i++ {
sp.appendText("ab")
time.Sleep(10 * time.Millisecond)
}
// Wait for throttle timers to fire
time.Sleep(300 * time.Millisecond)
msgs := mp.getMessages()
// Should NOT have 10 individual updates; throttling should batch them
if len(msgs) >= 10 {
t.Errorf("expected throttling to reduce updates, got %d", len(msgs))
}
if len(msgs) == 0 {
t.Error("expected at least one update")
}
}
func TestStreamPreview_MaxChars(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 50,
MinDeltaChars: 1,
MaxChars: 10,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
sp.appendText("This is a very long text that exceeds max chars limit")
time.Sleep(100 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) == 0 {
t.Fatal("expected at least one message")
}
// Last message should be truncated
for _, m := range msgs {
if len(m) > 0 {
// Content after "start:" or "update:" should respect maxChars
content := m
for _, prefix := range []string{"start:", "update:"} {
if len(content) > len(prefix) && content[:len(prefix)] == prefix {
content = content[len(prefix):]
}
}
if len([]rune(content)) > 15 { // 10 chars + "…" + "▍" with some margin
t.Errorf("message too long: %q (%d runes)", content, len([]rune(content)))
}
}
}
}
func TestStreamPreview_Disabled(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{Enabled: false}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
if sp.canPreview() {
t.Error("should not be able to preview when disabled")
}
sp.appendText("Hello")
time.Sleep(50 * time.Millisecond)
msgs := mp.getMessages()
if len(msgs) != 0 {
t.Error("no messages should be sent when disabled")
}
}
func TestStreamPreview_FinishInPlace(t *testing.T) {
mp := &mockUpdaterPlatform{}
cfg := StreamPreviewCfg{
Enabled: true,
IntervalMs: 50,
MinDeltaChars: 1,
MaxChars: 500,
}
sp := newStreamPreview(cfg, mp, "ctx", context.Background())
sp.appendText("Hello World")
time.Sleep(100 * time.Millisecond)
ok := sp.finish("Hello World Final")
if !ok {
t.Error("finish should return true when preview was active")
}
msgs := mp.getMessages()
last := msgs[len(msgs)-1]
if last != "update:Hello World Final" {
t.Errorf("last message = %q, want 'update:Hello World Final'", last)
}
}
func TestStreamPreview_NonUpdaterPlatform(t *testing.T) {
p := &stubPlatformEngine{n: "plain"}
cfg := DefaultStreamPreviewCfg()
sp := newStreamPreview(cfg, p, "ctx", context.Background())
if sp.canPreview() {
t.Error("should not preview on non-updater platform")
}
}