add WithReporter and WithAssertionHandler (#376)

Co-authored-by: Victor Gaydov <victor@enise.org>

chain.go

@@ -272,6 +272,20 @@ func (c *chain) setResponse(resp *Response) {
 	c.context.Response = resp
 }
 
+// Set assertion handler
+// Chain always overrides assertion handler with given one.
+func (c *chain) setHandler(handler AssertionHandler) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if chainValidation && c.state == stateLeaved {
+		panic("can't use chain after leave")
+	}
+
+	c.handler = handler
+	c.context.TestingTB = isTestingTB(handler)
+}
+
 // Create chain clone.
 // Typically is called between enter() and leave().
 func (c *chain) clone() *chain {

request.go

@@ -242,6 +242,94 @@ func (r *Request) WithName(name string) *Request {
 	return r
 }
 
+// WithReporter sets reporter to be used for this request.
+//
+// The new reporter overwrites AssertionHandler.
+// The new AssertionHandler is DefaultAssertionHandler with specified reporter,
+// existing Config.Formatter and nil Logger.
+// It will be used to report formatted fatal failure messages.
+//
+// Example:
+//
+//	req := NewRequestC(config, "GET", "http://example.com/path")
+//	req.WithReporter(t)
+func (r *Request) WithReporter(reporter Reporter) *Request {
+	opChain := r.chain.enter("WithReporter()")
+	defer opChain.leave()
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if opChain.failed() {
+		return r
+	}
+
+	if !r.checkOrder(opChain, "WithReporter()") {
+		return r
+	}
+
+	if reporter == nil {
+		opChain.fail(AssertionFailure{
+			Type: AssertUsage,
+			Errors: []error{
+				errors.New("unexpected nil argument"),
+			},
+		})
+		return r
+	}
+
+	handler := &DefaultAssertionHandler{
+		Reporter:  reporter,
+		Formatter: r.config.Formatter,
+	}
+	r.chain.setHandler(handler)
+
+	return r
+}
+
+// WithAssertionHandler sets assertion handler to be used for this request.
+//
+// The new handler overwrites assertion handler that will be used
+// by Request and its children (Response, Body, etc.).
+// It will be used to format and report test Failure or Success.
+//
+// Example:
+//
+//	req := NewRequestC(config, "GET", "http://example.com/path")
+//	req.WithAssertionHandler(&DefaultAssertionHandler{
+//		Reporter:  reporter,
+//		Formatter: formatter,
+//	})
+func (r *Request) WithAssertionHandler(handler AssertionHandler) *Request {
+	opChain := r.chain.enter("WithAssertionHandler()")
+	defer opChain.leave()
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if opChain.failed() {
+		return r
+	}
+
+	if !r.checkOrder(opChain, "WithAssertionHandler()") {
+		return r
+	}
+
+	if handler == nil {
+		opChain.fail(AssertionFailure{
+			Type: AssertUsage,
+			Errors: []error{
+				errors.New("unexpected nil argument"),
+			},
+		})
+		return r
+	}
+
+	r.chain.setHandler(handler)
+
+	return r
+}
+
 // WithMatcher attaches a matcher to the request.
 // All attached matchers are invoked in the Expect method for a newly
 // created Response.

request_test.go

@@ -24,6 +24,7 @@ import (
 
 func TestRequest_FailedChain(t *testing.T) {
 	reporter := newMockReporter(t)
+	assertionHandler := &mockAssertionHandler{}
 	config := newMockConfig(reporter)
 	chain := newChainWithDefaults("test", reporter, flagFailed)
 
@@ -32,6 +33,8 @@ func TestRequest_FailedChain(t *testing.T) {
 
 	req.Alias("foo")
 	req.WithName("foo")
+	req.WithReporter(reporter)
+	req.WithAssertionHandler(assertionHandler)
 	req.WithMatcher(func(resp *Response) {
 	})
 	req.WithTransformer(func(r *http.Request) {
@@ -267,6 +270,80 @@ func TestRequest_Basic(t *testing.T) {
 	})
 }
 
+func TestRequest_Reporter(t *testing.T) {
+	t.Run("default reporter", func(t *testing.T) {
+		configReporter := newMockReporter(t)
+		config := Config{
+			Client:   &mockClient{},
+			Reporter: configReporter,
+		}
+
+		req := NewRequestC(config, "GET", "/")
+
+		req.WithClient(nil)
+		req.chain.assert(t, failure)
+
+		assert.True(t, configReporter.reported)
+	})
+
+	t.Run("custom reporter", func(t *testing.T) {
+		configReporter := newMockReporter(t)
+		config := Config{
+			Client:   &mockClient{},
+			Reporter: configReporter,
+		}
+
+		req := NewRequestC(config, "GET", "/")
+
+		// set custom reporter
+		reqReporter := newMockReporter(t)
+		req.WithReporter(reqReporter)
+
+		req.WithClient(nil)
+		req.chain.assert(t, failure)
+
+		assert.False(t, configReporter.reported)
+		assert.True(t, reqReporter.reported)
+	})
+}
+
+func TestRequest_AssertionHandler(t *testing.T) {
+	t.Run("default handler", func(t *testing.T) {
+		configHandler := &mockAssertionHandler{}
+		config := Config{
+			Client:           &mockClient{},
+			AssertionHandler: configHandler,
+		}
+
+		req := NewRequestC(config, "GET", "/")
+
+		req.WithClient(nil)
+		req.chain.assert(t, failure)
+
+		assert.Equal(t, 1, configHandler.failureCalled)
+	})
+
+	t.Run("custom handler", func(t *testing.T) {
+		configHandler := &mockAssertionHandler{}
+		config := Config{
+			Client:           &mockClient{},
+			AssertionHandler: configHandler,
+		}
+
+		req := NewRequestC(config, "GET", "/")
+
+		// set custom assertion handler
+		reqHandler := &mockAssertionHandler{}
+		req.WithAssertionHandler(reqHandler)
+
+		req.WithClient(nil)
+		req.chain.assert(t, failure)
+
+		assert.Equal(t, 0, configHandler.failureCalled)
+		assert.Equal(t, 1, reqHandler.failureCalled)
+	})
+}
+
 func TestRequest_Matchers(t *testing.T) {
 	client := &mockClient{}
 
@@ -3426,6 +3503,22 @@ func TestRequest_Usage(t *testing.T) {
 		prepFails   bool
 		expectFails bool
 	}{
+		{
+			name: "WithReporter - nil argument",
+			prepFunc: func(req *Request) {
+				req.WithReporter(nil)
+			},
+			prepFails:   true,
+			expectFails: true,
+		},
+		{
+			name: "WithAssertionHandler - nil argument",
+			prepFunc: func(req *Request) {
+				req.WithAssertionHandler(nil)
+			},
+			prepFails:   true,
+			expectFails: true,
+		},
 		{
 			name: "WithMatcher - nil argument",
 			prepFunc: func(req *Request) {
@@ -3600,6 +3693,18 @@ func TestRequest_Order(t *testing.T) {
 				req.WithName("Test")
 			},
 		},
+		{
+			name: "WithReporter after Expect",
+			afterFunc: func(req *Request) {
+				req.WithReporter(newMockReporter(t))
+			},
+		},
+		{
+			name: "WithAssertionHandler after Expect",
+			afterFunc: func(req *Request) {
+				req.WithAssertionHandler(&mockAssertionHandler{})
+			},
+		},
 		{
 			name: "WithMatcher after Expect",
 			afterFunc: func(req *Request) {