learning-go-for-sre-interview/08-graceful-shutdown/main_test.go

package main

import (
	"context"
	"io"
	"net"
	"net/http"
	"testing"
	"time"
)

// testClient disables keep-alives: the default transport may speculatively
// dial a second connection and park it in StateNew without ever sending a
// request — and Server.Shutdown waits for StateNew conns, deadlocking the
// drain. A fun wart, but not the one this exercise is about.
var testClient = &http.Client{Transport: &http.Transport{DisableKeepAlives: true}}

func freeAddr(t *testing.T) string {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	addr := l.Addr().String()
	l.Close()
	return addr
}

func waitReady(t *testing.T, url string) {
	t.Helper()
	for i := 0; i < 100; i++ {
		if resp, err := testClient.Get(url); err == nil {
			resp.Body.Close()
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("server never became ready")
}

func TestServerHasTimeouts(t *testing.T) {
	srv := NewServer(":0", http.NewServeMux())
	if srv.ReadTimeout == 0 || srv.WriteTimeout == 0 || srv.IdleTimeout == 0 {
		t.Fatalf("server must set Read/Write/Idle timeouts explicitly (got %v/%v/%v) — zero means unlimited",
			srv.ReadTimeout, srv.WriteTimeout, srv.IdleTimeout)
	}
	if srv.Addr != ":0" || srv.Handler == nil {
		t.Fatal("Addr/Handler not wired through")
	}
}

func TestServesAndStopsOnCancel(t *testing.T) {
	addr := freeAddr(t)
	mux := http.NewServeMux()
	mux.HandleFunc("/ok", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "ok") })

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- Run(ctx, NewServer(addr, mux), 5*time.Second) }()

	waitReady(t, "http://"+addr+"/ok")
	cancel()

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("clean shutdown should return nil, got %v (leaking http.ErrServerClosed?)", err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("Run did not return after ctx cancel")
	}
}

func TestInFlightRequestCompletes(t *testing.T) {
	addr := freeAddr(t)
	started := make(chan struct{})
	mux := http.NewServeMux()
	mux.HandleFunc("/ok", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "ok") })
	mux.HandleFunc("/slow", func(w http.ResponseWriter, r *http.Request) {
		close(started)
		time.Sleep(500 * time.Millisecond)
		io.WriteString(w, "slow-done")
	})

	ctx, cancel := context.WithCancel(context.Background())
	runDone := make(chan error, 1)
	go func() { runDone <- Run(ctx, NewServer(addr, mux), 5*time.Second) }()
	waitReady(t, "http://"+addr+"/ok")

	reqDone := make(chan string, 1)
	go func() {
		resp, err := testClient.Get("http://" + addr + "/slow")
		if err != nil {
			reqDone <- "error: " + err.Error()
			return
		}
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		reqDone <- string(b)
	}()

	<-started // request is in flight
	cancel()  // SIGTERM arrives mid-request

	if body := <-reqDone; body != "slow-done" {
		t.Fatalf("in-flight request was killed instead of drained: got %q", body)
	}
	if err := <-runDone; err != nil {
		t.Fatalf("drain within deadline should return nil, got %v", err)
	}
}

func TestDrainDeadlineExceeded(t *testing.T) {
	addr := freeAddr(t)
	started := make(chan struct{})
	mux := http.NewServeMux()
	mux.HandleFunc("/ok", func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, "ok") })
	mux.HandleFunc("/hang", func(w http.ResponseWriter, r *http.Request) {
		close(started)
		time.Sleep(5 * time.Second)
	})

	ctx, cancel := context.WithCancel(context.Background())
	runDone := make(chan error, 1)
	go func() { runDone <- Run(ctx, NewServer(addr, mux), 200*time.Millisecond) }()
	waitReady(t, "http://"+addr+"/ok")

	go func() { testClient.Get("http://" + addr + "/hang") }()
	<-started
	cancel()

	select {
	case err := <-runDone:
		if err == nil {
			t.Fatal("want error when drain deadline exceeded, got nil")
		}
	case <-time.After(3 * time.Second):
		t.Fatal("Run hung past the drain deadline — Shutdown ctx wired up wrong?")
	}
}