package server_test import ( "encoding/json" "net/http" "net/http/httptest" "testing" "git.codelab.vc/pkg/httpx/server" ) func TestRouter_NotFoundHandler(t *testing.T) { t.Run("custom 404 handler", func(t *testing.T) { notFound := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusNotFound) json.NewEncoder(w).Encode(map[string]string{"error": "not found"}) }) r := server.NewRouter(server.WithNotFoundHandler(notFound)) r.HandleFunc("GET /exists", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }) // Matched route works normally. w := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/exists", nil) r.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Fatalf("matched route: got status %d, want %d", w.Code, http.StatusOK) } // Unmatched route uses custom handler. w = httptest.NewRecorder() req = httptest.NewRequest(http.MethodGet, "/nope", nil) r.ServeHTTP(w, req) if w.Code != http.StatusNotFound { t.Fatalf("not found: got status %d, want %d", w.Code, http.StatusNotFound) } if ct := w.Header().Get("Content-Type"); ct != "application/json" { t.Fatalf("Content-Type = %q, want %q", ct, "application/json") } var body map[string]string if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil { t.Fatalf("unmarshal: %v", err) } if body["error"] != "not found" { t.Fatalf("error = %q, want %q", body["error"], "not found") } }) t.Run("default behavior without custom handler", func(t *testing.T) { r := server.NewRouter() r.HandleFunc("GET /exists", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) }) w := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/nope", nil) r.ServeHTTP(w, req) // Default ServeMux returns 404. if w.Code != http.StatusNotFound { t.Fatalf("got status %d, want %d", w.Code, http.StatusNotFound) } }) }