httpx/server/server_test.go

package server_test

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

	"git.codelab.vc/pkg/httpx/server"
)

func TestServerLifecycle(t *testing.T) {
	t.Run("starts and serves requests", func(t *testing.T) {
		handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("hello"))
		})

		srv := server.New(handler, server.WithAddr(":0"))

		// Start in background and wait for addr.
		errCh := make(chan error, 1)
		go func() { errCh <- srv.ListenAndServe() }()

		waitForAddr(t, srv)

		resp, err := http.Get("http://" + srv.Addr())
		if err != nil {
			t.Fatalf("GET failed: %v", err)
		}
		defer resp.Body.Close()

		body, _ := io.ReadAll(resp.Body)
		if string(body) != "hello" {
			t.Fatalf("got body %q, want %q", body, "hello")
		}

		// Shutdown.
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := srv.Shutdown(ctx); err != nil {
			t.Fatalf("shutdown failed: %v", err)
		}
	})

	t.Run("addr returns empty before start", func(t *testing.T) {
		srv := server.New(http.NotFoundHandler())
		if addr := srv.Addr(); addr != "" {
			t.Fatalf("got addr %q before start, want empty", addr)
		}
	})
}

func TestGracefulShutdown(t *testing.T) {
	t.Run("calls onShutdown hooks", func(t *testing.T) {
		called := false
		handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
		})

		srv := server.New(handler,
			server.WithAddr(":0"),
			server.WithOnShutdown(func() { called = true }),
		)

		go func() { _ = srv.ListenAndServe() }()
		waitForAddr(t, srv)

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := srv.Shutdown(ctx); err != nil {
			t.Fatalf("shutdown failed: %v", err)
		}

		if !called {
			t.Fatal("onShutdown hook was not called")
		}
	})
}

func TestServerWithMiddleware(t *testing.T) {
	t.Run("applies middleware from options", func(t *testing.T) {
		var called bool
		mw := func(next http.Handler) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				called = true
				next.ServeHTTP(w, r)
			})
		}

		handler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
		})

		srv := server.New(handler,
			server.WithAddr(":0"),
			server.WithMiddleware(mw),
		)

		go func() { _ = srv.ListenAndServe() }()
		waitForAddr(t, srv)

		resp, err := http.Get("http://" + srv.Addr())
		if err != nil {
			t.Fatalf("GET failed: %v", err)
		}
		resp.Body.Close()

		if !called {
			t.Fatal("middleware was not called")
		}

		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = srv.Shutdown(ctx)
	})
}

// waitForAddr polls until the server's Addr() is non-empty.
func waitForAddr(t *testing.T, srv *server.Server) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if srv.Addr() != "" {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("server did not start in time")
}