Add ShouldCache option

Author: bep

httpcache.go

@@ -105,6 +105,9 @@ type Transport struct {
 	// a successful response from the cache will be returned without connecting to the server.
 	AlwaysUseCachedResponse func(req *http.Request, key string) bool
 
+	// ShouldCache is an optional func that when it returns false, the response will not be cached.
+	ShouldCache func(req *http.Request, resp *http.Response, key string) bool
+
 	// Around is an optional func.
 	// If set, the Transport will call Around at the start of RoundTrip
 	// and defer the returned func until the end of RoundTrip.
@@ -228,7 +231,7 @@ func (t *Transport) RoundTrip(req *http.Request) (resp *http.Response, err error
 		}
 	}
 
-	if cacheable && canStore(parseCacheControl(req.Header), parseCacheControl(resp.Header)) {
+	if cacheable && (t.ShouldCache == nil || t.ShouldCache(req, resp, cacheKey)) && canStore(parseCacheControl(req.Header), parseCacheControl(resp.Header)) {
 		for _, varyKey := range headerAllCommaSepValues(resp.Header, "vary") {
 			varyKey = http.CanonicalHeaderKey(varyKey)
 			fakeHeader := "X-Varied-" + varyKey

httpcache_test.go

@@ -178,6 +178,8 @@ func cacheSize() int {
 func resetTest() {
 	s.transport.Cache = newMemoryCache()
 	s.transport.CacheKey = nil
+	s.transport.AlwaysUseCachedResponse = nil
+	s.transport.ShouldCache = nil
 	s.transport.EnableETagPair = false
 	s.transport.MarkCachedResponses = false
 	clock = &realClock{}
@@ -269,6 +271,30 @@ func TestAlwaysUseCachedResponse(t *testing.T) {
 	}
 }
 
+func TestShouldCache(t *testing.T) {
+	resetTest()
+	c := qt.New(t)
+	s.transport.AlwaysUseCachedResponse = func(req *http.Request, key string) bool {
+		return true
+	}
+
+	s.transport.ShouldCache = func(req *http.Request, resp *http.Response, key string) bool {
+		return req.Header.Get("Hello") == "world2"
+	}
+	{
+		s, _ := doMethod(t, "GET", "/helloheaderasbody", map[string]string{"Hello": "world1"})
+		c.Assert(s, qt.Equals, "world1")
+	}
+	{
+		s, _ := doMethod(t, "GET", "/helloheaderasbody", map[string]string{"Hello": "world2"})
+		c.Assert(s, qt.Equals, "world2")
+	}
+	{
+		s, _ := doMethod(t, "GET", "/helloheaderasbody", map[string]string{"Hello": "world3"})
+		c.Assert(s, qt.Equals, "world2")
+	}
+}
+
 func TestAround(t *testing.T) {
 	resetTest()
 	c := qt.New(t)