From e358f9c1549197247bc02c41526206e18aa3c6cf Mon Sep 17 00:00:00 2001
From: Gregory Newman-Smith <gregory.newmansmith@couchbase.com>
Date: Fri, 3 Jan 2025 16:53:05 +0000
Subject: [PATCH] updates to address comments

---
 rest/api_test.go                     |  7 +------
 rest/counted_response_writer_test.go | 15 +++++++++++++++
 rest/handler.go                      |  9 ++++++++-
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/rest/api_test.go b/rest/api_test.go
index 598668c68b..60bfbf00e4 100644
--- a/rest/api_test.go
+++ b/rest/api_test.go
@@ -2873,17 +2873,12 @@ func TestBufferFlush(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, a.Save(user))
 
-	// create some changes
-	for i := 0; i < 10; i++ {
-		rt.PutDoc(fmt.Sprint(i), `{"some":"doc", "channels":["foo"]}`)
-	}
-
 	var wg sync.WaitGroup
 	var resp *TestResponse
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		resp = rt.SendUserRequest(http.MethodGet, "/{{.keyspace}}/_changes?feed=continuous&since=0&timeout=5000&include_docs=true", "", "foo")
+		resp = rt.SendUserRequest(http.MethodGet, "/{{.keyspace}}/_changes?feed=continuous&since=0&timeout=500&include_docs=true", "", "foo")
 		RequireStatus(t, resp, http.StatusOK)
 	}()
 	wg.Wait()
diff --git a/rest/counted_response_writer_test.go b/rest/counted_response_writer_test.go
index 79253c257c..50480d4ba6 100644
--- a/rest/counted_response_writer_test.go
+++ b/rest/counted_response_writer_test.go
@@ -16,6 +16,7 @@ import (
 
 	"github.com/couchbase/sync_gateway/base"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -169,3 +170,17 @@ func TestCountableResponseWriterWithDelay(t *testing.T) {
 	}
 
 }
+
+func TestResponseWriterSupportsFLush(t *testing.T) {
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+
+			stat, err := base.NewIntStat(base.SubsystemDatabaseKey, "http_bytes_written", base.StatUnitBytes, base.PublicRestBytesWrittenDesc, base.StatAddedVersion3dot1dot0, base.StatDeprecatedVersionNotDeprecated, base.StatStabilityCommitted, nil, nil, prometheus.CounterValue, 0)
+			require.NoError(t, err)
+			responseWriter := getResponseWriter(t, stat, test.name, 0)
+
+			_, ok := responseWriter.(http.Flusher)
+			assert.True(t, ok)
+		})
+	}
+}
diff --git a/rest/handler.go b/rest/handler.go
index f2a5a90530..6658f956f2 100644
--- a/rest/handler.go
+++ b/rest/handler.go
@@ -39,6 +39,13 @@ const (
 	minCompressibleJSONSize = 1000
 )
 
+var _ http.Flusher = &CountedResponseWriter{}
+var _ http.Flusher = &NonCountedResponseWriter{}
+var _ http.Flusher = &EncodedResponseWriter{}
+
+var _ http.Hijacker = &CountedResponseWriter{}
+var _ http.Hijacker = &NonCountedResponseWriter{}
+
 var ErrInvalidLogin = base.HTTPErrorf(http.StatusUnauthorized, "Invalid login")
 var ErrLoginRequired = base.HTTPErrorf(http.StatusUnauthorized, "Login required")
 
@@ -674,7 +681,7 @@ func (h *handler) validateAndWriteHeaders(method handlerMethod, accessPermission
 	// ensure wrapped ResponseWriter implements http.Flusher
 	_, ok := h.response.(http.Flusher)
 	if !ok {
-		return fmt.Errorf("ResponseWriter does not implement Flusher interface")
+		return fmt.Errorf("http.ResponseWriter %T does not implement Flusher interface", h.response)
 	}
 	return nil
 }