From 9b7e7bdce6e073d47635d57fbacfb3205cb3127e Mon Sep 17 00:00:00 2001
From: Dmitry Dorogin <dmirogin@ya.ru>
Date: Tue, 7 Aug 2018 01:44:32 +0300
Subject: [PATCH] Add tests for context.Stream (#1433)

---
 context_test.go | 36 ++++++++++++++++++++++++++++++++++++
 test_helpers.go | 21 +++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/context_test.go b/context_test.go
index 589ab336..40ae5bf8 100644
--- a/context_test.go
+++ b/context_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/gin-gonic/gin/binding"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/net/context"
+	"io"
 )
 
 var _ context.Context = &Context{}
@@ -1558,3 +1559,38 @@ func TestContextRenderDataFromReader(t *testing.T) {
 	assert.Equal(t, fmt.Sprintf("%d", contentLength), w.HeaderMap.Get("Content-Length"))
 	assert.Equal(t, extraHeaders["Content-Disposition"], w.HeaderMap.Get("Content-Disposition"))
 }
+
+func TestContextStream(t *testing.T) {
+	w := CreateTestResponseRecorder()
+	c, _ := CreateTestContext(w)
+
+	stopStream := true
+	c.Stream(func(w io.Writer) bool {
+		defer func() {
+			stopStream = false
+		}()
+
+		w.Write([]byte("test"))
+
+		return stopStream
+	})
+
+	assert.Equal(t, "testtest", w.Body.String())
+}
+
+func TestContextStreamWithClientGone(t *testing.T) {
+	w := CreateTestResponseRecorder()
+	c, _ := CreateTestContext(w)
+
+	c.Stream(func(writer io.Writer) bool {
+		defer func() {
+			w.closeClient()
+		}()
+
+		writer.Write([]byte("test"))
+
+		return true
+	})
+
+	assert.Equal(t, "test", w.Body.String())
+}
diff --git a/test_helpers.go b/test_helpers.go
index 2aed37f2..554568d9 100644
--- a/test_helpers.go
+++ b/test_helpers.go
@@ -6,6 +6,7 @@ package gin
 
 import (
 	"net/http"
+	"net/http/httptest"
 )
 
 // CreateTestContext returns a fresh engine and context for testing purposes
@@ -16,3 +17,23 @@ func CreateTestContext(w http.ResponseWriter) (c *Context, r *Engine) {
 	c.writermem.reset(w)
 	return
 }
+
+type TestResponseRecorder struct {
+	*httptest.ResponseRecorder
+	closeChannel chan bool
+}
+
+func (r *TestResponseRecorder) CloseNotify() <-chan bool {
+	return r.closeChannel
+}
+
+func (r *TestResponseRecorder) closeClient() {
+	r.closeChannel <- true
+}
+
+func CreateTestResponseRecorder() *TestResponseRecorder {
+	return &TestResponseRecorder{
+		httptest.NewRecorder(),
+		make(chan bool, 1),
+	}
+}