github.com/lianghucheng/zrddz@v0.0.0-20200923083010-c71f680932e2/src/golang.org/x/net/context/ctxhttp/ctxhttp_pre17_test.go (about)

     1  // Copyright 2015 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // +build !plan9,!go1.7
     6  
     7  package ctxhttp
     8  
     9  import (
    10  	"net"
    11  	"net/http"
    12  	"net/http/httptest"
    13  	"sync"
    14  	"testing"
    15  	"time"
    16  
    17  	"golang.org/x/net/context"
    18  )
    19  
    20  // golang.org/issue/14065
    21  func TestClosesResponseBodyOnCancel(t *testing.T) {
    22  	defer func() { testHookContextDoneBeforeHeaders = nop }()
    23  	defer func() { testHookDoReturned = nop }()
    24  	defer func() { testHookDidBodyClose = nop }()
    25  
    26  	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
    27  	defer ts.Close()
    28  
    29  	ctx, cancel := context.WithCancel(context.Background())
    30  
    31  	// closed when Do enters select case <-ctx.Done()
    32  	enteredDonePath := make(chan struct{})
    33  
    34  	testHookContextDoneBeforeHeaders = func() {
    35  		close(enteredDonePath)
    36  	}
    37  
    38  	testHookDoReturned = func() {
    39  		// We now have the result (the Flush'd headers) at least,
    40  		// so we can cancel the request.
    41  		cancel()
    42  
    43  		// But block the client.Do goroutine from sending
    44  		// until Do enters into the <-ctx.Done() path, since
    45  		// otherwise if both channels are readable, select
    46  		// picks a random one.
    47  		<-enteredDonePath
    48  	}
    49  
    50  	sawBodyClose := make(chan struct{})
    51  	testHookDidBodyClose = func() { close(sawBodyClose) }
    52  
    53  	tr := &http.Transport{}
    54  	defer tr.CloseIdleConnections()
    55  	c := &http.Client{Transport: tr}
    56  	req, _ := http.NewRequest("GET", ts.URL, nil)
    57  	_, doErr := Do(ctx, c, req)
    58  
    59  	select {
    60  	case <-sawBodyClose:
    61  	case <-time.After(5 * time.Second):
    62  		t.Fatal("timeout waiting for body to close")
    63  	}
    64  
    65  	if doErr != ctx.Err() {
    66  		t.Errorf("Do error = %v; want %v", doErr, ctx.Err())
    67  	}
    68  }
    69  
    70  type noteCloseConn struct {
    71  	net.Conn
    72  	onceClose sync.Once
    73  	closefn   func()
    74  }
    75  
    76  func (c *noteCloseConn) Close() error {
    77  	c.onceClose.Do(c.closefn)
    78  	return c.Conn.Close()
    79  }