nex/backend/internal/conversion/openai/stream_encoder_test.go

package openai

import (
	"encoding/json"
	"strings"
	"testing"

	"nex/backend/internal/conversion/canonical"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestStreamEncoder_MessageStart(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewMessageStartEvent("chatcmpl-1", "gpt-4")

	chunks := e.EncodeEvent(event)
	require.Len(t, chunks, 1)

	s := string(chunks[0])
	assert.True(t, strings.HasPrefix(s, "data: "))
	assert.Contains(t, s, "chatcmpl-1")
	assert.Contains(t, s, "chat.completion.chunk")

	var payload map[string]any
	data := strings.TrimPrefix(s, "data: ")
	data = strings.TrimRight(data, "\n")
	require.NoError(t, json.Unmarshal([]byte(data), &payload))
	choices := payload["choices"].([]any)
	delta := choices[0].(map[string]any)["delta"].(map[string]any)
	assert.Equal(t, "assistant", delta["role"])
}

func TestStreamEncoder_TextDelta(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{Type: "text_delta", Text: "你好"})

	chunks := e.EncodeEvent(event)
	require.Len(t, chunks, 1)

	s := string(chunks[0])
	assert.Contains(t, s, "你好")
}

func TestStreamEncoder_MessageStop(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewMessageStopEvent()

	chunks := e.EncodeEvent(event)
	require.Len(t, chunks, 1)
	assert.Equal(t, "data: [DONE]\n\n", string(chunks[0]))
}

func TestStreamEncoder_Buffering(t *testing.T) {
	e := NewStreamEncoder()

	// ContentBlockStart 应被缓冲，不输出
	startEvent := canonical.NewContentBlockStartEvent(0, canonical.StreamContentBlock{Type: "text", Text: ""})
	chunks := e.EncodeEvent(startEvent)
	assert.Nil(t, chunks)
	assert.NotNil(t, e.bufferedStart)

	// 第一个 delta 触发输出（清空缓冲）
	deltaEvent := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{Type: "text_delta", Text: "hello"})
	chunks = e.EncodeEvent(deltaEvent)
	require.NotEmpty(t, chunks)
	assert.Nil(t, e.bufferedStart)
}

func TestStreamEncoder_ContentBlockStop_ReturnsNil(t *testing.T) {
	e := NewStreamEncoder()
	idx := 0
	event := canonical.CanonicalStreamEvent{
		Type:  canonical.EventContentBlockStop,
		Index: &idx,
	}
	chunks := e.EncodeEvent(event)
	assert.Nil(t, chunks)
}

func TestStreamEncoder_Ping_ReturnsNil(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewPingEvent()
	chunks := e.EncodeEvent(event)
	assert.Nil(t, chunks)
}

func TestStreamEncoder_Error_ReturnsNil(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewErrorEvent("test_error", "测试错误")
	chunks := e.EncodeEvent(event)
	assert.Nil(t, chunks)
}

func TestStreamEncoder_Flush_ReturnsNil(t *testing.T) {
	e := NewStreamEncoder()
	chunks := e.Flush()
	assert.Nil(t, chunks)
}

func TestStreamEncoder_ThinkingDelta(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:     string(canonical.DeltaTypeThinking),
		Thinking: "思考内容",
	})

	chunks := e.EncodeEvent(event)
	require.Len(t, chunks, 1)

	s := string(chunks[0])
	assert.Contains(t, s, "reasoning_content")
	assert.Contains(t, s, "思考内容")
}

func TestStreamEncoder_InputJSONDelta(t *testing.T) {
	e := NewStreamEncoder()

	e.EncodeEvent(canonical.NewContentBlockStartEvent(0, canonical.StreamContentBlock{
		Type: "tool_use",
		ID:   "call_1",
		Name: "get_weather",
	}))

	event := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: "{\"city\":\"北京\"}",
	})

	chunks := e.EncodeEvent(event)
	require.NotEmpty(t, chunks)

	s := string(chunks[0])
	assert.Contains(t, s, "tool_calls")
	assert.Contains(t, s, "北京")
}

func TestStreamEncoder_MessageDelta_WithStopReason(t *testing.T) {
	e := NewStreamEncoder()
	sr := canonical.StopReasonEndTurn
	event := canonical.CanonicalStreamEvent{
		Type:       canonical.EventMessageDelta,
		StopReason: &sr,
	}

	chunks := e.EncodeEvent(event)
	require.NotEmpty(t, chunks)

	s := string(chunks[0])
	assert.Contains(t, s, "finish_reason")
	assert.Contains(t, s, "stop")
}

func TestStreamEncoder_MessageDelta_WithUsage(t *testing.T) {
	e := NewStreamEncoder()
	usage := canonical.CanonicalUsage{
		InputTokens:  100,
		OutputTokens: 50,
	}
	event := canonical.CanonicalStreamEvent{
		Type:  canonical.EventMessageDelta,
		Usage: &usage,
	}

	chunks := e.EncodeEvent(event)
	require.NotEmpty(t, chunks)

	s := string(chunks[0])
	assert.Contains(t, s, "usage")
	assert.Contains(t, s, "prompt_tokens")
}

func TestStreamEncoder_InputJSONDelta_SubsequentDelta(t *testing.T) {
	e := NewStreamEncoder()

	e.EncodeEvent(canonical.NewContentBlockStartEvent(0, canonical.StreamContentBlock{
		Type: "tool_use",
		ID:   "call_1",
		Name: "get_weather",
	}))

	e.EncodeEvent(canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: "{\"city\":",
	}))

	event := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: "\"Beijing\"}",
	})

	chunks := e.EncodeEvent(event)
	require.NotEmpty(t, chunks)

	s := string(chunks[0])
	assert.Contains(t, s, "tool_calls")
	assert.Contains(t, s, "Beijing")
}

func TestStreamEncoder_MessageStart_NilMessage(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.CanonicalStreamEvent{Type: canonical.EventMessageStart}

	chunks := e.EncodeEvent(event)
	require.Len(t, chunks, 1)

	s := string(chunks[0])
	assert.Contains(t, s, "chat.completion.chunk")
}

func TestStreamEncoder_UnknownEvent_ReturnsNil(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.CanonicalStreamEvent{Type: "unknown_type"}
	chunks := e.EncodeEvent(event)
	assert.Nil(t, chunks)
}

func TestStreamEncoder_ContentBlockDelta_NilDelta(t *testing.T) {
	e := NewStreamEncoder()
	event := canonical.CanonicalStreamEvent{Type: canonical.EventContentBlockDelta}
	chunks := e.EncodeEvent(event)
	assert.Nil(t, chunks)
}

func TestStreamEncoder_MultiToolCall_IndexMapping(t *testing.T) {
	e := NewStreamEncoder()

	e.EncodeEvent(canonical.NewContentBlockStartEvent(0, canonical.StreamContentBlock{
		Type: "tool_use",
		ID:   "call_1",
		Name: "get_weather",
	}))

	firstDelta := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: `{"city":"北京"}`,
	})
	chunks := e.EncodeEvent(firstDelta)
	require.NotEmpty(t, chunks)
	s := string(chunks[0])
	assert.Contains(t, s, `"index":0`)
	assert.Contains(t, s, "get_weather")
	assert.Contains(t, s, "北京")

	e.EncodeEvent(canonical.NewContentBlockStopEvent(0))

	e.EncodeEvent(canonical.NewContentBlockStartEvent(1, canonical.StreamContentBlock{
		Type: "tool_use",
		ID:   "call_2",
		Name: "get_time",
	}))

	secondDelta := canonical.NewContentBlockDeltaEvent(1, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: `{"tz":"Asia/Shanghai"}`,
	})
	chunks = e.EncodeEvent(secondDelta)
	require.NotEmpty(t, chunks)
	s = string(chunks[0])
	assert.Contains(t, s, `"index":1`)
	assert.Contains(t, s, "get_time")
	assert.Contains(t, s, "Asia/Shanghai")

	subsequentDelta0 := canonical.NewContentBlockDeltaEvent(0, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: `"more_data"`,
	})
	chunks = e.EncodeEvent(subsequentDelta0)
	require.NotEmpty(t, chunks)
	s = string(chunks[0])
	assert.Contains(t, s, `"index":0`)
	assert.NotContains(t, s, "get_weather")
	assert.Contains(t, s, "more_data")

	subsequentDelta1 := canonical.NewContentBlockDeltaEvent(1, canonical.StreamDelta{
		Type:        string(canonical.DeltaTypeInputJSON),
		PartialJSON: `"more_time"`,
	})
	chunks = e.EncodeEvent(subsequentDelta1)
	require.NotEmpty(t, chunks)
	s = string(chunks[0])
	assert.Contains(t, s, `"index":1`)
	assert.Contains(t, s, "more_time")
}