github.com/cavaliergopher/grab/v3@v3.0.1/pkg/grabtest/handler_option.go (about)

     1  package grabtest
     2  
     3  import (
     4  	"errors"
     5  	"net/http"
     6  	"time"
     7  )
     8  
     9  type HandlerOption func(*handler) error
    10  
    11  func StatusCodeStatic(code int) HandlerOption {
    12  	return func(h *handler) error {
    13  		return StatusCode(func(req *http.Request) int {
    14  			return code
    15  		})(h)
    16  	}
    17  }
    18  
    19  func StatusCode(f StatusCodeFunc) HandlerOption {
    20  	return func(h *handler) error {
    21  		if f == nil {
    22  			return errors.New("status code function cannot be nil")
    23  		}
    24  		h.statusCodeFunc = f
    25  		return nil
    26  	}
    27  }
    28  
    29  func MethodWhitelist(methods ...string) HandlerOption {
    30  	return func(h *handler) error {
    31  		h.methodWhitelist = methods
    32  		return nil
    33  	}
    34  }
    35  
    36  func HeaderBlacklist(headers ...string) HandlerOption {
    37  	return func(h *handler) error {
    38  		h.headerBlacklist = headers
    39  		return nil
    40  	}
    41  }
    42  
    43  func ContentLength(n int) HandlerOption {
    44  	return func(h *handler) error {
    45  		if n < 0 {
    46  			return errors.New("content length must be zero or greater")
    47  		}
    48  		h.contentLength = n
    49  		return nil
    50  	}
    51  }
    52  
    53  func AcceptRanges(enabled bool) HandlerOption {
    54  	return func(h *handler) error {
    55  		h.acceptRanges = enabled
    56  		return nil
    57  	}
    58  }
    59  
    60  func LastModified(t time.Time) HandlerOption {
    61  	return func(h *handler) error {
    62  		h.lastModified = t.UTC()
    63  		return nil
    64  	}
    65  }
    66  
    67  func TimeToFirstByte(d time.Duration) HandlerOption {
    68  	return func(h *handler) error {
    69  		if d < 1 {
    70  			return errors.New("time to first byte must be greater than zero")
    71  		}
    72  		h.ttfb = d
    73  		return nil
    74  	}
    75  }
    76  
    77  func RateLimiter(bps int) HandlerOption {
    78  	return func(h *handler) error {
    79  		if bps < 1 {
    80  			return errors.New("bytes per second must be greater than zero")
    81  		}
    82  		h.rateLimiter = time.NewTicker(time.Second / time.Duration(bps))
    83  		return nil
    84  	}
    85  }
    86  
    87  func AttachmentFilename(filename string) HandlerOption {
    88  	return func(h *handler) error {
    89  		h.attachmentFilename = filename
    90  		return nil
    91  	}
    92  }