Kevin Wan 2 роки тому
батько
коміт
83f88d177f
2 змінених файлів з 44 додано та 3 видалено
  1. 1 3
      rest/handler/tracinghandler.go
  2. 43 0
      rest/handler/tracinghandler_test.go

+ 1 - 3
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
 

+ 43 - 0
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() {
+}