@@ -810,6 +810,39 @@ func TestRetryableRequestBodyWithCloser(t *testing.T) {
810
810
require .True (t , tr .closeCalled )
811
811
}
812
812
813
+ func TestRetryPolicySuccessWithRetryPreserveHeaders (t * testing.T ) {
814
+ srv , close := mock .NewServer ()
815
+ defer close ()
816
+ srv .AppendResponse (mock .WithStatusCode (http .StatusRequestTimeout ))
817
+ srv .AppendResponse ()
818
+ pl := exported .NewPipeline (srv , NewRetryPolicy (testRetryOptions ()), exported .PolicyFunc (challengeLikePolicy ))
819
+ req , err := NewRequest (context .Background (), http .MethodGet , srv .URL ())
820
+ require .NoError (t , err )
821
+ body := newRewindTrackingBody ("stuff" )
822
+ require .NoError (t , req .SetBody (body , "text/plain" ))
823
+ resp , err := pl .Do (req )
824
+ require .NoError (t , err )
825
+ require .EqualValues (t , http .StatusOK , resp .StatusCode )
826
+ require .EqualValues (t , 2 , srv .Requests ())
827
+ require .EqualValues (t , 1 , body .rcount )
828
+ require .True (t , body .closed )
829
+ }
830
+
831
+ func challengeLikePolicy (req * policy.Request ) (* http.Response , error ) {
832
+ if req .Body () == nil {
833
+ return nil , errors .New ("request body wasn't restored" )
834
+ }
835
+ if req .Raw ().Header .Get ("content-type" ) != "text/plain" {
836
+ return nil , errors .New ("content-type header wasn't restored" )
837
+ }
838
+
839
+ // remove the body and header. the retry policy should restore them
840
+ if err := req .SetBody (nil , "" ); err != nil {
841
+ return nil , err
842
+ }
843
+ return req .Next ()
844
+ }
845
+
813
846
func newRewindTrackingBody (s string ) * rewindTrackingBody {
814
847
// there are two rewinds that happen before rewinding for a retry
815
848
// 1. to get the body's size in SetBody()
0 commit comments