vendor/golang.org/x/net/context/ctxhttp/ctxhttp_pre17_test.go

   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 }