github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/httpsnoop/wrap_test.go (about)

     1  package httpsnoop
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"io/ioutil"
     7  	"testing"
     8  
     9  	"github.com/hxx258456/ccgo/gmhttp/httptest"
    10  
    11  	http "github.com/hxx258456/ccgo/gmhttp"
    12  )
    13  
    14  func TestWrap_integration(t *testing.T) {
    15  	tests := []struct {
    16  		Name     string
    17  		Handler  http.Handler
    18  		Hooks    Hooks
    19  		WantCode int
    20  		WantBody []byte
    21  	}{
    22  		{
    23  			Name: "WriteHeader (no hook)",
    24  			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    25  				w.WriteHeader(http.StatusNotFound)
    26  			}),
    27  			WantCode: http.StatusNotFound,
    28  		},
    29  		{
    30  			Name: "WriteHeader",
    31  			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    32  				w.WriteHeader(http.StatusNotFound)
    33  			}),
    34  			Hooks: Hooks{
    35  				WriteHeader: func(next WriteHeaderFunc) WriteHeaderFunc {
    36  					return func(code int) {
    37  						if code != http.StatusNotFound {
    38  							t.Errorf("got=%d want=%d", code, http.StatusNotFound)
    39  						}
    40  						next(http.StatusForbidden)
    41  					}
    42  				},
    43  			},
    44  			WantCode: http.StatusForbidden,
    45  		},
    46  
    47  		{
    48  			Name: "Write (no hook)",
    49  			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    50  				w.Write([]byte("foo"))
    51  			}),
    52  			WantCode: http.StatusOK,
    53  			WantBody: []byte("foo"),
    54  		},
    55  		{
    56  			Name: "Write (rewrite hook)",
    57  			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    58  				if n, err := w.Write([]byte("foo")); err != nil {
    59  					t.Errorf("got=%s", err)
    60  				} else if got, want := n, len("foobar"); got != want {
    61  					t.Errorf("got=%d want=%d", got, want)
    62  				}
    63  			}),
    64  			Hooks: Hooks{
    65  				Write: func(next WriteFunc) WriteFunc {
    66  					return func(p []byte) (int, error) {
    67  						if string(p) != "foo" {
    68  							t.Errorf("%s", p)
    69  						}
    70  						return next([]byte("foobar"))
    71  					}
    72  				},
    73  			},
    74  			WantCode: http.StatusOK,
    75  			WantBody: []byte("foobar"),
    76  		},
    77  		{
    78  			Name: "Write (error hook)",
    79  			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    80  				if n, err := w.Write([]byte("foo")); n != 0 {
    81  					t.Errorf("got=%d want=%d", n, 0)
    82  				} else if err != io.EOF {
    83  					t.Errorf("got=%s want=%s", err, io.EOF)
    84  				}
    85  			}),
    86  			Hooks: Hooks{
    87  				Write: func(next WriteFunc) WriteFunc {
    88  					return func(p []byte) (int, error) {
    89  						if string(p) != "foo" {
    90  							t.Errorf("%s", p)
    91  						}
    92  						return 0, io.EOF
    93  					}
    94  				},
    95  			},
    96  			WantCode: http.StatusOK,
    97  		},
    98  	}
    99  
   100  	for _, test := range tests {
   101  		func() {
   102  			h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   103  				sw := Wrap(w, test.Hooks)
   104  				test.Handler.ServeHTTP(sw, r)
   105  			})
   106  			s := httptest.NewServer(h)
   107  			defer s.Close()
   108  			res, err := http.Get(s.URL)
   109  			if err != nil {
   110  				t.Fatal(err)
   111  			}
   112  			defer res.Body.Close()
   113  			gotBody, err := ioutil.ReadAll(res.Body)
   114  			if res.StatusCode != test.WantCode {
   115  				t.Errorf("got=%d want=%d", res.StatusCode, test.WantCode)
   116  			} else if !bytes.Equal(gotBody, test.WantBody) {
   117  				t.Errorf("got=%s want=%s", gotBody, test.WantBody)
   118  			}
   119  		}()
   120  	}
   121  }