Simplify injecting custom pipeline policies

sdk/armcore/connection.go

@@ -40,6 +40,14 @@ type ConnectionOptions struct {
 	// DisableRPRegistration disables the auto-RP registration policy.
 	// The default value is false.
 	DisableRPRegistration bool
+
+	// PerCallPolicies contains custom policies to inject into the pipeline.
+	// Each policy is executed once per request.
+	PerCallPolicies []azcore.Policy
+
+	// PerRetryPolicies contains custom policies to inject into the pipeline.
+	// Each policy is executed once per request, and for each retry request.
+	PerRetryPolicies []azcore.Policy
 }
 
 // Connection is a connection to an Azure Resource Manager endpoint.
@@ -83,16 +91,13 @@ func NewConnection(endpoint string, cred azcore.TokenCredential, options *Connec
 		}
 		policies = append(policies, NewRPRegistrationPolicy(endpoint, cred, &regRPOpts))
 	}
-	policies = append(policies, azcore.NewRetryPolicy(&options.Retry),
+	policies = append(policies, options.PerCallPolicies...)
+	policies = append(policies, azcore.NewRetryPolicy(&options.Retry))
+	policies = append(policies, options.PerRetryPolicies...)
+	policies = append(policies,
 		cred.AuthenticationPolicy(azcore.AuthenticationPolicyOptions{Options: azcore.TokenRequestOptions{Scopes: []string{endpointToScope(endpoint)}}}),
 		azcore.NewLogPolicy(&options.Logging))
 	p := azcore.NewPipeline(options.HTTPClient, policies...)
-	return NewConnectionWithPipeline(endpoint, p)
-}
-
-// NewConnectionWithPipeline creates an instance of the Connection type with the specified endpoint and pipeline.
-// Use this when a custom pipeline is required.
-func NewConnectionWithPipeline(endpoint string, p azcore.Pipeline) *Connection {
 	return &Connection{u: endpoint, p: p}
 }
 

sdk/armcore/connection_test.go

@@ -104,25 +104,6 @@ func TestNewConnectionWithCustomTelemetry(t *testing.T) {
 	}
 }
 
-func TestNewConnectionWithPipeline(t *testing.T) {
-	srv, close := mock.NewServer()
-	defer close()
-	srv.AppendResponse()
-	p := getPipeline(srv)
-	con := NewConnectionWithPipeline(srv.URL(), p)
-	req, err := azcore.NewRequest(context.Background(), http.MethodGet, srv.URL())
-	if err != nil {
-		t.Fatalf("Unexpected error: %v", err)
-	}
-	resp, err := con.Pipeline().Do(req)
-	if err != nil {
-		t.Fatalf("Unexpected error: %v", err)
-	}
-	if resp.StatusCode != http.StatusOK {
-		t.Fatalf("unexpected status code: %d", resp.StatusCode)
-	}
-}
-
 func TestScope(t *testing.T) {
 	if s := endpointToScope(AzureGermany); s != "https://management.microsoftazure.de//.default" {
 		t.Fatalf("unexpected scope %s", s)
@@ -167,3 +148,45 @@ func TestDisableAutoRPRegistration(t *testing.T) {
 		t.Fatalf("expected 0 log entries, got %d", logEntries)
 	}
 }
+
+// policy that tracks the number of times it was invoked
+type countingPolicy struct {
+	count int
+}
+
+func (p *countingPolicy) Do(req *azcore.Request) (*azcore.Response, error) {
+	p.count++
+	return req.Next()
+}
+
+func TestConnectionWithCustomPolicies(t *testing.T) {
+	srv, close := mock.NewServer()
+	defer close()
+	// initial response is a failure to trigger retry
+	srv.AppendResponse(mock.WithStatusCode(http.StatusInternalServerError))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK))
+	perCallPolicy := countingPolicy{}
+	perRetryPolicy := countingPolicy{}
+	con := NewConnection(srv.URL(), mockTokenCred{}, &ConnectionOptions{
+		DisableRPRegistration: true,
+		PerCallPolicies:       []azcore.Policy{&perCallPolicy},
+		PerRetryPolicies:      []azcore.Policy{&perRetryPolicy},
+	})
+	req, err := azcore.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp, err := con.Pipeline().Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("unexpected status code %d", resp.StatusCode)
+	}
+	if perCallPolicy.count != 1 {
+		t.Fatalf("unexpected per call policy count %d", perCallPolicy.count)
+	}
+	if perRetryPolicy.count != 2 {
+		t.Fatalf("unexpected per retry policy count %d", perRetryPolicy.count)
+	}
+}

sdk/armcore/version.go

@@ -10,5 +10,5 @@ const (
 	UserAgent = "armcore/" + Version
 
 	// Version is the semantic version (see http://semver.org) of this module.
-	Version = "v0.5.1"
+	Version = "v0.7.0"
 )