From 83f88d177fce8ea03924241dbfe22fc8c9caac0c Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Wed, 4 Jan 2023 13:42:20 +0800
Subject: [PATCH] chore: improve codecov (#2752)

---
 rest/handler/tracinghandler.go      |  4 +--
 rest/handler/tracinghandler_test.go | 43 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go
index 7d6b285f..43ad9d2c 100644
--- a/rest/handler/tracinghandler.go
+++ b/rest/handler/tracinghandler.go
@@ -17,9 +17,7 @@ import (
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
-const (
-	traceKeyStatusCode = "http.status_code"
-)
+const traceKeyStatusCode = "http.status_code"
 
 var notTracingSpans sync.Map
 
diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go
index 79a94c4d..d4aebe2a 100644
--- a/rest/handler/tracinghandler_test.go
+++ b/rest/handler/tracinghandler_test.go
@@ -148,3 +148,46 @@ func TestTraceResponseWriter(t *testing.T) {
 		})
 	}
 }
+
+func TestTraceHandler_traceResponseWriter(t *testing.T) {
+	writer := &traceResponseWriter{
+		w: httptest.NewRecorder(),
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &traceResponseWriter{
+		w: mockedHijackable{httptest.NewRecorder()},
+	}
+	assert.NotPanics(t, func() {
+		writer.Hijack()
+	})
+
+	writer = &traceResponseWriter{
+		w: httptest.NewRecorder(),
+	}
+	writer.WriteHeader(http.StatusBadRequest)
+	assert.NotNil(t, writer.Header())
+
+	writer = &traceResponseWriter{
+		w: httptest.NewRecorder(),
+	}
+	assert.NotPanics(t, func() {
+		writer.Flush()
+	})
+
+	writer = &traceResponseWriter{
+		w: mockedFlusher{httptest.NewRecorder()},
+	}
+	assert.NotPanics(t, func() {
+		writer.Flush()
+	})
+}
+
+type mockedFlusher struct {
+	*httptest.ResponseRecorder
+}
+
+func (m mockedFlusher) Flush() {
+}