Polling works for custom action -- need to fix conversion to the expe…

…cted response
radius-project · Dec 15, 2022 · 63a4a20 · 63a4a20
1 parent 68edd73
commit 63a4a20
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/pkg/azure/clientv2/customaction.go b/pkg/azure/clientv2/customaction.go
@@ -43,6 +43,7 @@ func NewCustomActionClientWithBaseURI(baseURI string, subscriptionID string, cre
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},
@@ -70,8 +71,12 @@ func NewCustomActionClientWithBaseURI(baseURI string, subscriptionID string, cre
 }
 
 type ClientCustomActionResponse struct {
-	Body     *map[string]any
-	Response *http.Response
+	CustomActionResponseBody
+}
+
+type CustomActionResponseBody struct {
+	Body     *map[string]any `json:"body,omitempty"`
+	Response *http.Response  `json:"response,omitempty"`
 }
 
 type ClientBeginCustomActionOptions struct {
@@ -82,7 +87,6 @@ type ClientBeginCustomActionOptions struct {
 
 // NewClientCustomActionOptions creates an instance of the CustomActionClientOptions.
 func NewClientBeginCustomActionOptions(resourceID, action, apiVersion string) *ClientBeginCustomActionOptions {
-	// FIXME: This is to validate the resourceID.
 	_, err := resources.ParseResource(resourceID)
 	if err != nil {
 		return nil
@@ -101,7 +105,10 @@ func (client *CustomActionClient) BeginCustomAction(ctx context.Context, opts *C
 		return nil, err
 	}
 
-	return runtime.NewPoller[ClientCustomActionResponse](resp, *client.Pipeline, nil)
+	return runtime.NewPoller(resp, *client.Pipeline, &runtime.NewPollerOptions[ClientCustomActionResponse]{
+		FinalStateVia: runtime.FinalStateViaAzureAsyncOp,
+		Response:      &ClientCustomActionResponse{},
+	})
 }
 
 func (client *CustomActionClient) customAction(ctx context.Context, opts *ClientBeginCustomActionOptions) (*http.Response, error) {
@@ -114,14 +121,17 @@ func (client *CustomActionClient) customAction(ctx context.Context, opts *Client
 	if err != nil {
 		return nil, err
 	}
-	if !runtime.HasStatusCode(resp, http.StatusAccepted, http.StatusNoContent) {
+
+	if !runtime.HasStatusCode(resp, http.StatusOK, http.StatusAccepted, http.StatusNoContent) {
 		return nil, runtime.NewResponseError(resp)
 	}
+
 	return resp, nil
 }
 
 func (client *CustomActionClient) customActionCreateRequest(ctx context.Context, opts *ClientBeginCustomActionOptions) (*policy.Request, error) {
 	urlPath := "/{resourceID}/{action}"
+
 	if opts.resourceID == "" {
 		return nil, errors.New("resourceID cannot be empty")
 	}
@@ -132,7 +142,6 @@ func (client *CustomActionClient) customActionCreateRequest(ctx context.Context,
 	}
 	urlPath = strings.ReplaceAll(urlPath, "{action}", url.PathEscape(opts.action))
 
-	// FIXME: Is joining BaseURI and URLPath going to give us a wrong URL?
 	req, err := runtime.NewRequest(ctx, http.MethodPost, runtime.JoinPaths(client.BaseURI, urlPath))
 	if err != nil {
 		return nil, err

diff --git a/pkg/azure/clientv2/resourcedeploymentclient.go b/pkg/azure/clientv2/resourcedeploymentclient.go
@@ -111,6 +111,7 @@ func NewDeploymentsClientWithBaseURI(credential azcore.TokenCredential, subscrip
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},

diff --git a/pkg/azure/clientv2/resourcedeploymentoperationsclient.go b/pkg/azure/clientv2/resourcedeploymentoperationsclient.go
@@ -39,6 +39,7 @@ func NewResourceDeploymentOperationsClientWithBaseURI(cred azcore.TokenCredentia
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},