diff --git a/examples_test.go b/examples_test.go index 94568265..8bdd66b0 100644 --- a/examples_test.go +++ b/examples_test.go @@ -518,7 +518,9 @@ func Example_hijack_requests() { // Not calling this will require you to mock the entire response. // This can be done with the SetXxx (Status, Header, Body) functions on the // ctx.Response struct. - _ = ctx.LoadResponse(http.DefaultClient, true) + _ = ctx.LoadResponse(http.DefaultClient, func(res *http.Response) bool { + return true + }) // Here we append some code to every js file. // The code will update the document title to "hi" diff --git a/hijack.go b/hijack.go index a02d4f33..7f8aa310 100644 --- a/hijack.go +++ b/hijack.go @@ -221,7 +221,7 @@ func (h *Hijack) ContinueRequest(cq *proto.FetchContinueRequest) { } // LoadResponse will send request to the real destination and load the response as default response to override. -func (h *Hijack) LoadResponse(client *http.Client, loadBody bool) error { +func (h *Hijack) LoadResponse(client *http.Client, shouldLoadBody func(*http.Response) bool) error { res, err := client.Do(h.Request.req) if err != nil { return err @@ -238,7 +238,7 @@ func (h *Hijack) LoadResponse(client *http.Client, loadBody bool) error { } } - if loadBody { + if shouldLoadBody(res) { b, err := io.ReadAll(res.Body) if err != nil { return err @@ -255,6 +255,11 @@ type HijackRequest struct { req *http.Request } +// Event associated to the request. +func (ctx *HijackRequest) Event() *proto.FetchRequestPaused { + return ctx.event +} + // Type of the resource. func (ctx *HijackRequest) Type() proto.NetworkResourceType { return ctx.event.ResourceType diff --git a/hijack_test.go b/hijack_test.go index 470fd8ce..67941c59 100644 --- a/hijack_test.go +++ b/hijack_test.go @@ -294,14 +294,18 @@ func TestHijackLoadResponseErr(t *testing.T) { router.MustAdd("http://localhost/a", func(ctx *rod.Hijack) { g.Err(ctx.LoadResponse(&http.Client{ Transport: &MockRoundTripper{err: errors.New("err")}, - }, true)) + }, func(*http.Response) bool { + return true + })) g.Err(ctx.LoadResponse(&http.Client{ Transport: &MockRoundTripper{res: &http.Response{ StatusCode: http.StatusOK, Body: io.NopCloser(&MockReader{err: errors.New("err")}), }}, - }, true)) + }, func(*http.Response) bool { + return true + })) wg.Done() diff --git a/must.go b/must.go index fa940210..4920ced6 100644 --- a/must.go +++ b/must.go @@ -1139,7 +1139,9 @@ func (r *HijackRouter) MustStop() { // MustLoadResponse is similar to [Hijack.LoadResponse]. func (h *Hijack) MustLoadResponse() { - h.browser.e(h.LoadResponse(http.DefaultClient, true)) + h.browser.e(h.LoadResponse(http.DefaultClient, func(res *http.Response) bool { + return true + })) } // MustEqual is similar to [Element.Equal].