github.com/jpmicrosoft/grab/v3@v3.0.2/pkg/grabtest/handler_test.go (about)

     1  package grabtest
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"net/http"
     7  	"testing"
     8  	"time"
     9  )
    10  
    11  func TestHandlerDefaults(t *testing.T) {
    12  	WithTestServer(t, func(url string) {
    13  		resp := MustHTTPDo(MustHTTPNewRequest("GET", url, nil))
    14  		AssertHTTPResponseStatusCode(t, resp, http.StatusOK)
    15  		AssertHTTPResponseContentLength(t, resp, 1048576)
    16  		AssertHTTPResponseHeader(t, resp, "Accept-Ranges", "bytes")
    17  	})
    18  }
    19  
    20  func TestHandlerMethodWhitelist(t *testing.T) {
    21  	tests := []struct {
    22  		Whitelist        []string
    23  		Method           string
    24  		ExpectStatusCode int
    25  	}{
    26  		{[]string{"GET", "HEAD"}, "GET", http.StatusOK},
    27  		{[]string{"GET", "HEAD"}, "HEAD", http.StatusOK},
    28  		{[]string{"GET"}, "HEAD", http.StatusMethodNotAllowed},
    29  		{[]string{"HEAD"}, "GET", http.StatusMethodNotAllowed},
    30  	}
    31  
    32  	for _, test := range tests {
    33  		WithTestServer(t, func(url string) {
    34  			resp := MustHTTPDoWithClose(MustHTTPNewRequest(test.Method, url, nil))
    35  			AssertHTTPResponseStatusCode(t, resp, test.ExpectStatusCode)
    36  		}, MethodWhitelist(test.Whitelist...))
    37  	}
    38  }
    39  
    40  func TestHandlerHeaderBlacklist(t *testing.T) {
    41  	contentLength := 4096
    42  	WithTestServer(t, func(url string) {
    43  		resp := MustHTTPDo(MustHTTPNewRequest("GET", url, nil))
    44  		defer resp.Body.Close()
    45  		if resp.ContentLength != -1 {
    46  			t.Errorf("expected Response.ContentLength: -1, got: %d", resp.ContentLength)
    47  		}
    48  		AssertHTTPResponseHeader(t, resp, "Content-Length", "")
    49  		AssertHTTPResponseBodyLength(t, resp, int64(contentLength))
    50  	},
    51  		ContentLength(contentLength),
    52  		HeaderBlacklist("Content-Length"),
    53  	)
    54  }
    55  
    56  func TestHandlerStatusCodeFuncs(t *testing.T) {
    57  	expect := 418 // I'm a teapot
    58  	WithTestServer(t, func(url string) {
    59  		resp := MustHTTPDo(MustHTTPNewRequest("GET", url, nil))
    60  		defer resp.Body.Close()
    61  		AssertHTTPResponseStatusCode(t, resp, expect)
    62  	},
    63  		StatusCode(func(req *http.Request) int { return expect }),
    64  	)
    65  }
    66  
    67  func TestHandlerContentLength(t *testing.T) {
    68  	tests := []struct {
    69  		Method          string
    70  		ContentLength   int
    71  		ExpectHeaderLen int64
    72  		ExpectBodyLen   int
    73  	}{
    74  		{"GET", 321, 321, 321},
    75  		{"HEAD", 321, 321, 0},
    76  		{"GET", 0, 0, 0},
    77  		{"HEAD", 0, 0, 0},
    78  	}
    79  
    80  	for _, test := range tests {
    81  		WithTestServer(t, func(url string) {
    82  			resp := MustHTTPDo(MustHTTPNewRequest(test.Method, url, nil))
    83  			defer resp.Body.Close()
    84  
    85  			AssertHTTPResponseHeader(t, resp, "Content-Length", "%d", test.ExpectHeaderLen)
    86  
    87  			b, err := ioutil.ReadAll(resp.Body)
    88  			if err != nil {
    89  				panic(err)
    90  			}
    91  			if len(b) != test.ExpectBodyLen {
    92  				t.Errorf(
    93  					"expected body length: %v, got: %v, in: %v",
    94  					test.ExpectBodyLen,
    95  					len(b),
    96  					test,
    97  				)
    98  			}
    99  		},
   100  			ContentLength(test.ContentLength),
   101  		)
   102  	}
   103  }
   104  
   105  func TestHandlerAcceptRanges(t *testing.T) {
   106  	header := "Accept-Ranges"
   107  	n := 128
   108  	t.Run("Enabled", func(t *testing.T) {
   109  		WithTestServer(t, func(url string) {
   110  			req := MustHTTPNewRequest("GET", url, nil)
   111  			req.Header.Set("Range", fmt.Sprintf("bytes=%d-", n/2))
   112  			resp := MustHTTPDo(req)
   113  			AssertHTTPResponseHeader(t, resp, header, "bytes")
   114  			AssertHTTPResponseContentLength(t, resp, int64(n/2))
   115  		},
   116  			ContentLength(n),
   117  		)
   118  	})
   119  
   120  	t.Run("Disabled", func(t *testing.T) {
   121  		WithTestServer(t, func(url string) {
   122  			req := MustHTTPNewRequest("GET", url, nil)
   123  			req.Header.Set("Range", fmt.Sprintf("bytes=%d-", n/2))
   124  			resp := MustHTTPDo(req)
   125  			AssertHTTPResponseHeader(t, resp, header, "")
   126  			AssertHTTPResponseContentLength(t, resp, int64(n))
   127  		},
   128  			AcceptRanges(false),
   129  			ContentLength(n),
   130  		)
   131  	})
   132  }
   133  
   134  func TestHandlerAttachmentFilename(t *testing.T) {
   135  	filename := "foo.pdf"
   136  	WithTestServer(t, func(url string) {
   137  		resp := MustHTTPDoWithClose(MustHTTPNewRequest("GET", url, nil))
   138  		AssertHTTPResponseHeader(t, resp, "Content-Disposition", `attachment;filename="%s"`, filename)
   139  	},
   140  		AttachmentFilename(filename),
   141  	)
   142  }
   143  
   144  func TestHandlerLastModified(t *testing.T) {
   145  	WithTestServer(t, func(url string) {
   146  		resp := MustHTTPDoWithClose(MustHTTPNewRequest("GET", url, nil))
   147  		AssertHTTPResponseHeader(t, resp, "Last-Modified", "Thu, 29 Nov 1973 21:33:09 GMT")
   148  	},
   149  		LastModified(time.Unix(123456789, 0)),
   150  	)
   151  }