Repository: labstack/echo
File: middleware/proxy_1_11_test.go (Go)

Code before:
	// must add targets:
	for _, target := range targets {
		assert.True(t, rb.AddTarget(target))
	}

	// must ignore duplicates:
	for _, target := range targets {
		assert.False(t, rb.AddTarget(target))
	}

	// Random
	e := echo.New()
	e.Use(Proxy(rb))
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	rec := httptest.NewRecorder()

	// Remote unreachable
	rec = httptest.NewRecorder()
	req.URL.Path = "/api/users"
	e.ServeHTTP(rec, req)
	assert.Equal(t, "/api/users", req.URL.Path)
	assert.Equal(t, http.StatusBadGateway, rec.Code)

	// client closed connection
	HTTPTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(5 * time.Second)
		w.WriteHeader(http.StatusOK)
	}))
	defer HTTPTarget.Close()
	targetURL, _ := url.Parse(HTTPTarget.URL)
	target := &ProxyTarget{
		Name: "target",
		URL:  targetURL,
	}
	rb = NewRandomBalancer(nil)
	assert.True(t, rb.AddTarget(target))
	e = echo.New()
	e.Use(Proxy(rb))
	rec = httptest.NewRecorder()
	req = httptest.NewRequest(http.MethodGet, "/", nil)
	ctx, cancel := context.WithCancel(req.Context())
	req = req.WithContext(ctx)
	go func() {
		startTime := time.Now()
		time.Sleep(2 * time.Second)
		cancel()
		endTime := time.Now()
		assert.Less(t, endTime.Sub(startTime).Seconds(), float64(5))
	}()
	e.ServeHTTP(rec, req)
	assert.Equal(t, 499, rec.Code)

Review comment (performance):
Using time.sleep with seconds makes test execution very slow (5sec)

You can replace with test case with this:
```go
func TestClientCancelConnectionResultsHTTPCode499(t *testing.T) {
	timeoutStopChan := make(chan struct{}, 0)
	HTTPTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-timeoutStopChan // wait until we have asserted timeout
		w.WriteHeader(http.StatusOK)
	}))
	defer HTTPTarget.Close()
	targetURL, _ := url.Parse(HTTPTarget.URL)
	target := &ProxyTarget{
		Name: "target",
		URL:  targetURL,
	}
	rb := NewRandomBalancer(nil)
	assert.True(t, rb.AddTarget(target))

	e := echo.New()
	e.Use(Proxy(rb))
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	ctx, cancel := context.WithCancel(req.Context())
	req = req.WithContext(ctx)

	go func() {
		startTime := time.Now()
		time.Sleep(10 * time.Millisecond)
		cancel()
		endTime := time.Now()
		assert.Less(t, endTime.Sub(startTime).Seconds(), float64(5))
	}()

	e.ServeHTTP(rec, req)

	assert.Equal(t, 499, rec.Code)
	timeoutStopChan <- struct{}{} // signal that handler can return and not block ending of this test
}
```

Code after:
	rb := NewRandomBalancer(nil)
	// must add targets:
	for _, target := range targets {
		assert.True(t, rb.AddTarget(target))
	}

	// must ignore duplicates:
	for _, target := range targets {
		assert.False(t, rb.AddTarget(target))
	}

	// Random
	e := echo.New()
	e.Use(Proxy(rb))
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	rec := httptest.NewRecorder()

	// Remote unreachable
	rec = httptest.NewRecorder()
	req.URL.Path = "/api/users"
	e.ServeHTTP(rec, req)
	assert.Equal(t, "/api/users", req.URL.Path)
	assert.Equal(t, http.StatusBadGateway, rec.Code)
}

func TestClientCancelConnectionResultsHTTPCode499(t *testing.T) {
	var timeoutStop sync.WaitGroup
	timeoutStop.Add(1)
	HTTPTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		timeoutStop.Wait() // wait until we have canceled the request
		w.WriteHeader(http.StatusOK)
	}))
	defer HTTPTarget.Close()
	targetURL, _ := url.Parse(HTTPTarget.URL)
	target := &ProxyTarget{
		Name: "target",
		URL:  targetURL,
	}
	rb := NewRandomBalancer(nil)
	assert.True(t, rb.AddTarget(target))
	e := echo.New()
	e.Use(Proxy(rb))
	rec := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	ctx, cancel := context.WithCancel(req.Context())
	req = req.WithContext(ctx)
	go func() {
		time.Sleep(10 * time.Millisecond)
		cancel()
	}()
	e.ServeHTTP(rec, req)