Add const for TTFB timeout

2 files changed, 9 insertions, 5 deletions

diff --git a/internal/httptransport/transport.go b/internal/httptransport/transport.go
index fd5473c8..d8e6a3fe 100644
--- a/internal/httptransport/transport.go
+++ b/internal/httptransport/transport.go
@@ -16,6 +16,13 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const (
+	// DefaultTTFBTimeout is the timeout used in the meteredRoundTripper
+	// when calling http.Transport.RoundTrip. The request will be cancelled
+	// if the response takes longer than this.
+	DefaultTTFBTimeout = 15 * time.Second
+)
+
 var (
 	sysPoolOnce = &sync.Once{}
 	sysPool     *x509.CertPool
@@ -98,9 +105,7 @@ func (mrt *meteredRoundTripper) RoundTrip(r *http.Request) (*http.Response, erro
 	ctx := httptrace.WithClientTrace(r.Context(), mrt.newTracer(start))
 	ctx, cancel := context.WithCancel(ctx)
 
-	timer := time.AfterFunc(mrt.ttfbTimeout, func() {
-		cancel()
-	})
+	timer := time.AfterFunc(mrt.ttfbTimeout, cancel)
 	defer timer.Stop()
 
 	r = r.WithContext(ctx)
diff --git a/internal/httptransport/transport_test.go b/internal/httptransport/transport_test.go
index 869f21db..47167a75 100644
--- a/internal/httptransport/transport_test.go
+++ b/internal/httptransport/transport_test.go
@@ -115,11 +115,10 @@ type mockRoundTripper struct {
 }
 
 func (mrt *mockRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
-	time.Sleep(mrt.timeout)
 	select {
 	case <-r.Context().Done():
 		return nil, r.Context().Err()
-	default:
+	case <-time.After(mrt.timeout):
 		return mrt.res, mrt.err
 	}
 }