Repository: go-chi/chi
File: mux_test.go (Go)

Code before:
	if _, body := testRequest(t, ts, "GET", "/admin2/nope", nil); body != "root 404" {
		t.Fatalf(body)
	}
}

func TestMuxWith(t *testing.T) {
	var cmwInit1, cmwHandler1 uint64
	var cmwInit2, cmwHandler2 uint64
	mw1 := func(next http.Handler) http.Handler {
		cmwInit1++
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			cmwHandler1++
			r = r.WithContext(context.WithValue(r.Context(), "inline1", "yes"))
			next.ServeHTTP(w, r)
		})
	}
	mw2 := func(next http.Handler) http.Handler {
		cmwInit2++
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			cmwHandler2++
			r = r.WithContext(context.WithValue(r.Context(), "inline2", "yes"))
			next.ServeHTTP(w, r)
		})
	}

	r := NewRouter()
	r.Get("/hi", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("bye"))
	})
	r.With(mw1).With(mw2).Get("/inline", func(w http.ResponseWriter, r *http.Request) {
		v1 := r.Context().Value("inline1").(string)
		v2 := r.Context().Value("inline2").(string)
		w.Write([]byte(fmt.Sprintf("inline %s %s", v1, v2)))
	})

	ts := httptest.NewServer(r)
	defer ts.Close()

	if _, body := testRequest(t, ts, "GET", "/hi", nil); body != "bye" {
		t.Fatalf(body)
	}
	if _, body := testRequest(t, ts, "GET", "/inline", nil); body != "inline yes yes" {
		t.Fatalf(body)
	}
	if cmwInit1 != 1 {
		t.Fatalf("expecting cmwInit1 to be 1, got %d", cmwInit1)
	}
	if cmwHandler1 != 1 {
		t.Fatalf("expecting cmwHandler1 to be 1, got %d", cmwHandler1)
	}
	if cmwInit2 != 1 {

Review comment (bug):
yea, I do like the idea that NotFound would be used all the way down from this point in the tree. If another subrouter defined NotFound it would use that one all the way down. Nice work.

However, what would happen in this test func if r.NotFound() was defined at the very top before the Mount(), I think it might fail


Code after:
	if _, body := testRequest(t, ts, "GET", "/admin2/nope", nil); body != "root 404" {
		t.Fatalf(body)
	}
}

func TestMuxComplicatedNotFound(t *testing.T) {
	// sub router with groups
	sub := NewRouter()
	sub.Route("/resource", func(r Router) {
		r.Get("/", func(w http.ResponseWriter, r *http.Request) {
			w.Write([]byte("private get"))
		})
	})

	// Root router with groups
	r := NewRouter()
	r.Get("/auth", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("auth get"))
	})
	r.Route("/public", func(r Router) {
		r.Get("/", func(w http.ResponseWriter, r *http.Request) {
			w.Write([]byte("public get"))
		})
	})
	r.Mount("/private", sub)
	r.NotFound(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("custom not-found"))
	})

	ts := httptest.NewServer(r)
	defer ts.Close()

	// check that we didn't broke correct routes
	if _, body := testRequest(t, ts, "GET", "/auth", nil); body != "auth get" {
		t.Fatalf(body)
	}
	if _, body := testRequest(t, ts, "GET", "/public", nil); body != "public get" {
		t.Fatalf(body)
	}
	if _, body := testRequest(t, ts, "GET", "/private/resource", nil); body != "private get" {
		t.Fatalf(body)
	}
	// check custom not-found on all levels
	if _, body := testRequest(t, ts, "GET", "/nope", nil); body != "custom not-found" {
		t.Fatalf(body)
	}
	if _, body := testRequest(t, ts, "GET", "/public/nope", nil); body != "custom not-found" {
		t.Fatalf(body)
	}
	if _, body := testRequest(t, ts, "GET", "/private/nope", nil); body != "custom not-found" {
		t.Fatalf(body)