Add support for custom PrepareDecorator chains

CHANGELOG.md

@@ -1,5 +1,11 @@
 # CHANGELOG
 
+## v12.4.0
+
+### New Features
+
+- Added `autorest.WithPrepareDecorators` and `autorest.GetPrepareDecorators` for adding and retrieving a custom chain of PrepareDecorators to the provided context.
+
 ## v12.3.0
 
 ### New Features

autorest/azure/async.go

@@ -417,6 +417,11 @@ func (pt *pollingTrackerBase) pollForStatus(ctx context.Context, sender autorest
 	}
 
 	req = req.WithContext(ctx)
+	preparer := autorest.CreatePreparer(autorest.GetPrepareDecorators(ctx)...)
+	req, err = preparer.Prepare(req)
+	if err != nil {
+		return autorest.NewErrorWithError(err, "pollingTrackerBase", "pollForStatus", nil, "failed preparing HTTP request")
+	}
 	pt.resp, err = sender.Do(req)
 	if err != nil {
 		return autorest.NewErrorWithError(err, "pollingTrackerBase", "pollForStatus", nil, "failed to send HTTP request")

autorest/preparer.go

@@ -16,6 +16,7 @@ package autorest
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"encoding/xml"
 	"fmt"
@@ -38,6 +39,27 @@ const (
 	headerUserAgent        = "User-Agent"
 )
 
+// used as a key type in context.WithValue()
+type ctxPrepareDecorators struct{}
+
+// WithPrepareDecorators adds the specified PrepareDecorators to the provided context.
+// If no PrepareDecorators are provided the context is unchanged.
+func WithPrepareDecorators(ctx context.Context, prepareDecorator []PrepareDecorator) context.Context {
+	if len(prepareDecorator) == 0 {
+		return ctx
+	}
+	return context.WithValue(ctx, ctxPrepareDecorators{}, prepareDecorator)
+}
+
+// GetPrepareDecorators returns the PrepareDecorators in the provided context or the provided default PrepareDecorators.
+func GetPrepareDecorators(ctx context.Context, defaultPrepareDecorators ...PrepareDecorator) []PrepareDecorator {
+	inCtx := ctx.Value(ctxPrepareDecorators{})
+	if pd, ok := inCtx.([]PrepareDecorator); ok {
+		return pd
+	}
+	return defaultPrepareDecorators
+}
+
 // Preparer is the interface that wraps the Prepare method.
 //
 // Prepare accepts and possibly modifies an http.Request (e.g., adding Headers). Implementations

autorest/preparer_test.go

@@ -15,6 +15,7 @@ package autorest
 //  limitations under the License.
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -851,3 +852,26 @@ func TestModifyingRequestWithExistingQueryParameters(t *testing.T) {
 		t.Fatalf("autorest: Preparing an existing request failed when setting the query parameters (%s)", r.URL.RawQuery)
 	}
 }
+
+func TestGetPrepareDecorators(t *testing.T) {
+	pd := GetPrepareDecorators(context.Background())
+	if l := len(pd); l != 0 {
+		t.Fatalf("expected zero length but got %d", l)
+	}
+	pd = GetPrepareDecorators(context.Background(), WithNothing(), AsFormURLEncoded())
+	if l := len(pd); l != 2 {
+		t.Fatalf("expected length of two but got %d", l)
+	}
+}
+
+func TestWithPrepareDecorators(t *testing.T) {
+	ctx := WithPrepareDecorators(context.Background(), []PrepareDecorator{WithUserAgent("somestring")})
+	pd := GetPrepareDecorators(ctx)
+	if l := len(pd); l != 1 {
+		t.Fatalf("expected length of one but got %d", l)
+	}
+	pd = GetPrepareDecorators(ctx, WithNothing(), WithNothing())
+	if l := len(pd); l != 1 {
+		t.Fatalf("expected length of one but got %d", l)
+	}
+}

autorest/sender.go

@@ -30,7 +30,11 @@ import (
 type ctxSendDecorators struct{}
 
 // WithSendDecorators adds the specified SendDecorators to the provided context.
+// If no SendDecorators are provided the context is unchanged.
 func WithSendDecorators(ctx context.Context, sendDecorator []SendDecorator) context.Context {
+	if len(sendDecorator) == 0 {
+		return ctx
+	}
 	return context.WithValue(ctx, ctxSendDecorators{}, sendDecorator)
 }
 

autorest/version.go

@@ -19,7 +19,7 @@ import (
 	"runtime"
 )
 
-const number = "v12.3.0"
+const number = "v12.4.0"
 
 var (
 	userAgent = fmt.Sprintf("Go/%s (%s-%s) go-autorest/%s",