File size: 2,209 Bytes
6bc074c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | package httpstream
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
)
func setupTestContext() (*gin.Context, *httptest.ResponseRecorder) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request, _ = http.NewRequest("POST", "/test", nil)
return c, w
}
func TestWriteSSEHeader(t *testing.T) {
c, w := setupTestContext()
WriteSSEHeader(c)
if w.Header().Get("Content-Type") != "text/event-stream" {
t.Errorf("expected Content-Type text/event-stream, got %s", w.Header().Get("Content-Type"))
}
if w.Header().Get("Cache-Control") != "no-cache" {
t.Errorf("expected Cache-Control no-cache, got %s", w.Header().Get("Cache-Control"))
}
if w.Header().Get("Connection") != "keep-alive" {
t.Errorf("expected Connection keep-alive, got %s", w.Header().Get("Connection"))
}
}
func TestWriteSSEEvent(t *testing.T) {
c, w := setupTestContext()
payload := map[string]string{"key": "value"}
ok := WriteSSEEvent(c, "test.event", payload)
if !ok {
t.Error("WriteSSEEvent should return true")
}
body := w.Body.String()
if body == "" {
t.Error("expected non-empty response body")
}
}
func TestWriteDone(t *testing.T) {
c, w := setupTestContext()
ok := WriteDone(c)
if !ok {
t.Error("WriteDone should return true")
}
body := w.Body.String()
if body != "data: [DONE]\n\n" {
t.Errorf("expected 'data: [DONE]\\n\\n', got %q", body)
}
}
func TestWriteImageStreamHeader(t *testing.T) {
c, w := setupTestContext()
WriteImageStreamHeader(c)
if w.Header().Get("Content-Type") != "text/event-stream" {
t.Errorf("expected Content-Type text/event-stream, got %s", w.Header().Get("Content-Type"))
}
}
func TestWriteImageStreamDone(t *testing.T) {
c, w := setupTestContext()
ok := WriteImageStreamDone(c)
if !ok {
t.Error("WriteImageStreamDone should return true")
}
body := w.Body.String()
if body != "data: [DONE]\n\n" {
t.Errorf("expected 'data: [DONE]\\n\\n', got %q", body)
}
}
func TestWriteImageStreamError(t *testing.T) {
c, w := setupTestContext()
WriteImageStreamError(c, 0, 1, "test error")
body := w.Body.String()
if body == "" {
t.Error("expected non-empty response body")
}
}
|