From 56995a2b8d5268b117aaeecb28022cd7b13b2e07 Mon Sep 17 00:00:00 2001
From: adw1n <adw1n@users.noreply.github.com>
Date: Mon, 21 Sep 2020 21:23:29 +0200
Subject: [PATCH 1/2] contrib/go-chi/chi: Handle 0 status code correctly

---
 contrib/go-chi/chi/chi.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/contrib/go-chi/chi/chi.go b/contrib/go-chi/chi/chi.go
index f09329c8e5..3f2f29c90d 100644
--- a/contrib/go-chi/chi/chi.go
+++ b/contrib/go-chi/chi/chi.go
@@ -61,6 +61,10 @@ func Middleware(opts ...Option) func(next http.Handler) http.Handler {
 
 			// set the status code
 			status := ww.Status()
+			// 0 status means one has not yet been sent in which case net/http library will write StatusOK
+			if ww.Status() == 0 {
+				status = http.StatusOK
+			}
 			span.SetTag(ext.HTTPCode, strconv.Itoa(status))
 
 			if status >= 500 && status < 600 {

From 105dcad31684ba658e59f89fd828ad88fb05bee4 Mon Sep 17 00:00:00 2001
From: adw1n <adw1n@users.noreply.github.com>
Date: Mon, 28 Sep 2020 22:19:51 +0200
Subject: [PATCH 2/2] Add no response written test

---
 contrib/go-chi/chi/chi_test.go | 85 +++++++++++++++++++++-------------
 1 file changed, 53 insertions(+), 32 deletions(-)

diff --git a/contrib/go-chi/chi/chi_test.go b/contrib/go-chi/chi/chi_test.go
index 42afdae987..27a1328585 100644
--- a/contrib/go-chi/chi/chi_test.go
+++ b/contrib/go-chi/chi/chi_test.go
@@ -39,42 +39,63 @@ func TestChildSpan(t *testing.T) {
 }
 
 func TestTrace200(t *testing.T) {
-	assert := assert.New(t)
-	mt := mocktracer.Start()
-	defer mt.Stop()
+	assertDoRequest := func(assert *assert.Assertions, mt mocktracer.Tracer, router *chi.Mux) {
+		r := httptest.NewRequest("GET", "/user/123", nil)
+		w := httptest.NewRecorder()
 
-	router := chi.NewRouter()
-	router.Use(Middleware(WithServiceName("foobar")))
-	router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
-		span, ok := tracer.SpanFromContext(r.Context())
-		assert.True(ok)
-		assert.Equal(span.(mocktracer.Span).Tag(ext.ServiceName), "foobar")
-		id := chi.URLParam(r, "id")
-		w.Write([]byte(id))
-	})
+		// do and verify the request
+		router.ServeHTTP(w, r)
+		response := w.Result()
+		assert.Equal(response.StatusCode, 200)
 
-	r := httptest.NewRequest("GET", "/user/123", nil)
-	w := httptest.NewRecorder()
+		// verify traces look good
+		spans := mt.FinishedSpans()
+		assert.Len(spans, 1)
+		if len(spans) < 1 {
+			t.Fatalf("no spans")
+		}
+		span := spans[0]
+		assert.Equal("http.request", span.OperationName())
+		assert.Equal(ext.SpanTypeWeb, span.Tag(ext.SpanType))
+		assert.Equal("foobar", span.Tag(ext.ServiceName))
+		assert.Equal("GET /user/{id}", span.Tag(ext.ResourceName))
+		assert.Equal("200", span.Tag(ext.HTTPCode))
+		assert.Equal("GET", span.Tag(ext.HTTPMethod))
+		assert.Equal("/user/123", span.Tag(ext.HTTPURL))
+	}
 
-	// do and verify the request
-	router.ServeHTTP(w, r)
-	response := w.Result()
-	assert.Equal(response.StatusCode, 200)
+	t.Run("response written", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
 
-	// verify traces look good
-	spans := mt.FinishedSpans()
-	assert.Len(spans, 1)
-	if len(spans) < 1 {
-		t.Fatalf("no spans")
-	}
-	span := spans[0]
-	assert.Equal("http.request", span.OperationName())
-	assert.Equal(ext.SpanTypeWeb, span.Tag(ext.SpanType))
-	assert.Equal("foobar", span.Tag(ext.ServiceName))
-	assert.Equal("GET /user/{id}", span.Tag(ext.ResourceName))
-	assert.Equal("200", span.Tag(ext.HTTPCode))
-	assert.Equal("GET", span.Tag(ext.HTTPMethod))
-	assert.Equal("/user/123", span.Tag(ext.HTTPURL))
+		router := chi.NewRouter()
+		router.Use(Middleware(WithServiceName("foobar")))
+		router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+			span, ok := tracer.SpanFromContext(r.Context())
+			assert.True(ok)
+			assert.Equal(span.(mocktracer.Span).Tag(ext.ServiceName), "foobar")
+			id := chi.URLParam(r, "id")
+			_, err := w.Write([]byte(id))
+			assert.NoError(err)
+		})
+		assertDoRequest(assert, mt, router)
+	})
+
+	t.Run("no response written", func(t *testing.T) {
+		assert := assert.New(t)
+		mt := mocktracer.Start()
+		defer mt.Stop()
+
+		router := chi.NewRouter()
+		router.Use(Middleware(WithServiceName("foobar")))
+		router.Get("/user/{id}", func(w http.ResponseWriter, r *http.Request) {
+			span, ok := tracer.SpanFromContext(r.Context())
+			assert.True(ok)
+			assert.Equal(span.(mocktracer.Span).Tag(ext.ServiceName), "foobar")
+		})
+		assertDoRequest(assert, mt, router)
+	})
 }
 
 func TestError(t *testing.T) {