Testing HTTP caching in Go

main.go

package main

import (
	"bytes"
	"flag"
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	"log"
	"net/http"
	"strconv"
	"strings"
)

var root = flag.String("root", ".", "file system path")

func main() {
	http.HandleFunc("/black/", blackHandler)
	http.Handle("/", http.FileServer(http.Dir(*root)))
	log.Println("Listening on 8080")
	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal("ListenAndServe:", err)
	}
}

func blackHandler(w http.ResponseWriter, r *http.Request) {

	key := "black"
	e := `"` + key + `"`
	w.Header().Set("Etag", e)
	w.Header().Set("Cache-Control", "max-age=2592000") // 30 days

	if match := r.Header.Get("If-None-Match"); match != "" {
		if strings.Contains(match, e) {
			w.WriteHeader(http.StatusNotModified)
			return
		}
	}

	m := image.NewRGBA(image.Rect(0, 0, 240, 240))
	black := color.RGBA{0, 0, 0, 255}
	draw.Draw(m, m.Bounds(), &image.Uniform{black}, image.ZP, draw.Src)

	var img image.Image = m
	writeImage(w, &img)
}

// writeImage encodes an image 'img' in jpeg format and writes it into ResponseWriter.
func writeImage(w http.ResponseWriter, img *image.Image) {

	buffer := new(bytes.Buffer)
	if err := jpeg.Encode(buffer, *img, nil); err != nil {
		log.Println("unable to encode image.")
	}

	w.Header().Set("Content-Type", "image/jpeg")
	w.Header().Set("Content-Length", strconv.Itoa(len(buffer.Bytes())))
	if _, err := w.Write(buffer.Bytes()); err != nil {
		log.Println("unable to write image.")
	}
}

main_test.go

package main

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestBlackHandler(t *testing.T) {

	blackHandlerFunc := http.HandlerFunc(blackHandler)
	var etag string

	// first request
	if r, err := http.NewRequest("GET", "", nil); err != nil {
		t.Errorf("%v", err)
	} else {
		recorder := httptest.NewRecorder()
		blackHandlerFunc.ServeHTTP(recorder, r)
		if recorder.Code != http.StatusOK {
			t.Errorf("returned %v. Expected %v.", recorder.Code, http.StatusOK)
		}
		// record etag to test cache
		etag = recorder.Header().Get("Etag")
	}

	// test caching
	if r, err := http.NewRequest("GET", "", nil); err != nil {
		t.Errorf("%v", err)
	} else {
		r.Header.Set("If-None-Match", etag)
		recorder := httptest.NewRecorder()
		blackHandlerFunc.ServeHTTP(recorder, r)
		if recorder.Code != http.StatusNotModified {
			t.Errorf("returned %v. Expected %v.", recorder.Code, http.StatusNotModified)
		}
	}
}