add ReadAsSSEWithContext

dara/stream.go

@@ -3,6 +3,7 @@ package dara
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -10,7 +11,6 @@ import (
 	"strings"
 )
 
-// 定义 Event 结构体
 type SSEEvent struct {
 	Id    *string
 	Event *string
@@ -18,7 +18,6 @@ type SSEEvent struct {
 	Retry *int
 }
 
-// 解析单个事件
 func parseEvent(lines []string) *SSEEvent {
 	event := &SSEEvent{}
 	for _, line := range lines {
@@ -153,3 +152,64 @@ func ReadAsSSE(body io.ReadCloser, eventChannel chan *SSEEvent, errorChannel cha
 		}
 	}()
 }
+
+func ReadAsSSEWithContext(ctx context.Context, body io.ReadCloser, eventChannel chan *SSEEvent, errorChannel chan error) {
+
+	go func() {
+		defer func() {
+			body.Close()
+			close(eventChannel)
+		}()
+
+		reader := bufio.NewReader(body)
+		var eventLines []string
+
+		for {
+			select {
+			case <-ctx.Done():
+				errorChannel <- ctx.Err()
+				return
+			default:
+			}
+
+			line, err := reader.ReadString('\n')
+			if err != nil {
+				if err == io.EOF {
+					// Handle the end of the stream and possibly pending event
+					if len(eventLines) > 0 {
+						event := parseEvent(eventLines)
+						select {
+						case eventChannel <- event:
+						case <-ctx.Done():
+							errorChannel <- ctx.Err()
+							return
+						}
+					}
+					errorChannel <- nil
+					return
+				}
+				errorChannel <- err
+				return
+			}
+
+			line = strings.TrimRight(line, "\n")
+
+			if line == "" {
+				// End of an SSE event
+				if len(eventLines) > 0 {
+					event := parseEvent(eventLines)
+					select {
+					case eventChannel <- event:
+					case <-ctx.Done():
+						errorChannel <- ctx.Err()
+						return
+					}
+					eventLines = []string{} // Reset for the next event
+				}
+				continue
+			}
+
+			eventLines = append(eventLines, line)
+		}
+	}()
+}

dara/stream_test.go

@@ -1,9 +1,11 @@
 package dara
 
 import (
+	"context"
 	"io/ioutil"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/alibabacloud-go/tea/utils"
 )
@@ -342,3 +344,168 @@ func Test_parseEvent(t *testing.T) {
 		utils.AssertNil(t, event.Retry)
 	})
 }
+
+func Test_ReadAsSSEWithContext(t *testing.T) {
+	// Test case 1: Basic SSE event with context
+	t.Run("BasicSSEEventWithContext", func(t *testing.T) {
+		sseData := "data: hello world\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx := context.Background()
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		event := <-eventChannel
+		err := <-errorChannel
+
+		utils.AssertNil(t, err)
+		utils.AssertNotNil(t, event)
+		utils.AssertNotNil(t, event.Data)
+		utils.AssertEqual(t, "hello world", *event.Data)
+		utils.AssertNil(t, event.Event)
+		utils.AssertNil(t, event.Id)
+		utils.AssertNil(t, event.Retry)
+	})
+
+	// Test case 2: SSE event with context timeout
+	t.Run("SSEWithContextTimeout", func(t *testing.T) {
+		// 使用一个快过期的context来测试超时检查
+		sseData := "data: hello world\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Nanosecond)
+		defer cancel()
+
+		// 等待context超时
+		time.Sleep(10 * time.Millisecond)
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		// 应该立即收到超时错误
+		select {
+		case err := <-errorChannel:
+			utils.AssertNotNil(t, err)
+			utils.AssertEqual(t, context.DeadlineExceeded, err)
+		case <-time.After(50 * time.Millisecond):
+			t.Fatal("Expected timeout error but didn't receive one")
+		}
+	})
+
+	// Test case 3: SSE event with context cancellation
+	t.Run("SSEWithContextCancellation", func(t *testing.T) {
+		// 使用一个已经取消的context
+		sseData := "data: hello world\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel() // 立即取消
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		// 应该立即收到取消错误
+		select {
+		case err := <-errorChannel:
+			utils.AssertNotNil(t, err)
+			utils.AssertEqual(t, context.Canceled, err)
+		case <-time.After(50 * time.Millisecond):
+			t.Fatal("Expected cancellation error but didn't receive one")
+		}
+	})
+
+	// Test case 4: Complete SSE event with context
+	t.Run("CompleteSSEEventWithContext", func(t *testing.T) {
+		sseData := "id: 456\nevent: notification\ndata: welcome\ndata: to sse\nretry: 3000\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx := context.Background()
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		event := <-eventChannel
+		err := <-errorChannel
+
+		utils.AssertNil(t, err)
+		utils.AssertNotNil(t, event)
+		utils.AssertNotNil(t, event.Data)
+		utils.AssertEqual(t, "welcome\nto sse", *event.Data)
+		utils.AssertNotNil(t, event.Event)
+		utils.AssertEqual(t, "notification", *event.Event)
+		utils.AssertNotNil(t, event.Id)
+		utils.AssertEqual(t, "456", *event.Id)
+		utils.AssertNotNil(t, event.Retry)
+		utils.AssertEqual(t, 3000, *event.Retry)
+	})
+
+	// Test case 5: Multiple SSE events with context
+	t.Run("MultipleSSEEventsWithContext", func(t *testing.T) {
+		sseData := "data: first\n\n" + "data: second\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx := context.Background()
+
+		eventChannel := make(chan *SSEEvent, 2)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		event1 := <-eventChannel
+		event2 := <-eventChannel
+		err := <-errorChannel
+
+		utils.AssertNil(t, err)
+		utils.AssertNotNil(t, event1)
+		utils.AssertNotNil(t, event1.Data)
+		utils.AssertEqual(t, "first", *event1.Data)
+		utils.AssertNotNil(t, event2)
+		utils.AssertNotNil(t, event2.Data)
+		utils.AssertEqual(t, "second", *event2.Data)
+	})
+
+	// Test case 6: Context cancellation during event processing
+	t.Run("ContextCancellationDuringProcessing", func(t *testing.T) {
+		// 验证context取消机制的基本功能
+		sseData := "data: event1\n\n"
+		reader := ioutil.NopCloser(strings.NewReader(sseData))
+		ctx, cancel := context.WithCancel(context.Background())
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		// 读取第一个事件
+		event1 := <-eventChannel
+		utils.AssertNotNil(t, event1)
+		utils.AssertNotNil(t, event1.Data)
+		utils.AssertEqual(t, "event1", *event1.Data)
+
+		// 正常结束，收到nil错误
+		err := <-errorChannel
+		utils.AssertNil(t, err)
+
+		cancel() // 清理资源
+	})
+
+	// Test case 7: Empty stream with context
+	t.Run("EmptyStreamWithContext", func(t *testing.T) {
+		reader := ioutil.NopCloser(strings.NewReader(""))
+		ctx := context.Background()
+
+		eventChannel := make(chan *SSEEvent, 1)
+		errorChannel := make(chan error, 1)
+		ReadAsSSEWithContext(ctx, reader, eventChannel, errorChannel)
+
+		// 应该立即收到nil错误（EOF）
+		err := <-errorChannel
+		utils.AssertNil(t, err)
+
+		// 事件通道应该为空
+		select {
+		case event := <-eventChannel:
+			utils.AssertNil(t, event) // 不应该收到任何事件
+		default:
+			// 这是期望的行为
+		}
+	})
+}