Final response handler for Pollers and updating tests (#400)

* Removing operation response handler and adding custom handler for final response * Updating some lro tests with new design * Adding response handler as method on pollers, adding response handler method on pollingTracker, minor code fixes * Removing unnecessary test code * Improve tests, check payload * Adding special FinalResponse case for HTTPPollers * Fix HTTP pollers final response method and related tests * Removing unused handleresponse for http pollers and fixed PollUntilDone loop * Adding updated lro tests * Updating internal pollUntilDone to be a method on the unexported poller * Improving ResumeWrongPoller test * Removing time.Duration declarations in tests Co-authored-by: Catalina Peralta <[email protected]>
Azure · May 21, 2020 · 03099c8 · 03099c8
1 parent 6d8ed32
commit 03099c8
Show file tree

Hide file tree

Showing 11 changed files with 1,370 additions and 1,557 deletions.
diff --git a/src/generator/operations.ts b/src/generator/operations.ts
@@ -299,13 +299,18 @@ function generateOperation(clientName: string, op: Operation, imports: ImportMan
     text += `\tif err != nil {\n`;
     text += `\t\treturn nil, err\n`;
     text += `\t}\n`;
-    text += `\tresult.Poller = &${camelCase(op.language.go!.pollerType.name)}{\n`;
+    text += `\tpoller := &${camelCase(op.language.go!.pollerType.name)}{\n`;
     text += `\t\t\tpt: pt,\n`;
     text += `\t\t\tpipeline: client.p,\n`;
-    text += `\t\t\tresponse: client.${info.protocolNaming.responseMethod},\n`;
     text += `\t}\n`;
-    text += `\tresult.PollUntilDone = func(ctx context.Context, frequency time.Duration)(*${op.language.go!.pollerType.responseType}Response, error) {\n`;
-    text += `\t\treturn ${camelCase(op.language.go!.pollerType.name)}PollUntilDone(ctx, result.Poller, frequency)\n`;
+    text += '\tresult.Poller = poller\n';
+    // http pollers will simply return an *http.Response
+    if (op.language.go!.pollerType.name === 'HTTPPoller') {
+      text += `\tresult.PollUntilDone = func(ctx context.Context, frequency time.Duration) (*http.Response, error) {\n`;
+    } else {
+      text += `\tresult.PollUntilDone = func(ctx context.Context, frequency time.Duration)(*${op.language.go!.pollerType.responseType}Response, error) {\n`;
+    }
+    text += `\t\treturn poller.pollUntilDone(ctx, frequency)\n`;
     text += `\t}\n`;
     text += `\treturn result, nil\n`;
     // closing braces
@@ -954,7 +959,6 @@ function addResumePollerMethod(op: Operation, clientName: string): string {
   text += `\treturn &${camelCase(op.language.go!.pollerType.name)}{\n`;
   text += '\t\tpipeline: client.p,\n';
   text += '\t\tpt: pt,\n';
-  text += `\t\tresponse: client.${info.protocolNaming.responseMethod},\n`;
   text += '\t}, nil\n';
   text += '}\n';
   return text;

diff --git a/src/generator/pollers.ts b/src/generator/pollers.ts
@@ -34,35 +34,79 @@ export async function generatePollers(session: Session<CodeModel>): Promise<stri
     const pollerInterface = poller.name;
     const pollerName = camelCase(poller.name);
     let responseType = 'HTTPResponse';
-    let responseHandler = 'httpHandleResponse';
+    // HTTP Pollers do not need to perform the final get request since they do not return a model
+    let finalResponseDeclaration = 'FinalResponse() *http.Response';
+    let finalResponse = `${finalResponseDeclaration} {
+      return p.pt.latestResponse().Response;
+    }`;
+    let pollUntilDoneResponse = '(*http.Response, error)';
+    let pollUntilDoneReturn = 'p.FinalResponse(), nil';
+    let handleResponse = '';
     let rawResponse = ''; // used to access the raw response field on response envelopes
     const schemaResponse = <SchemaResponse>poller.op.responses![0];
     let unmarshalResponse = 'nil';
     if (isSchemaResponse(schemaResponse) && schemaResponse.schema.language.go!.responseType.value != undefined) {
-      unmarshalResponse = `resp.UnmarshalAsJSON(&result.${schemaResponse.schema.language.go!.responseType.value})`;
-    }
-    if (isSchemaResponse(schemaResponse)) {
-      responseHandler = `${camelCase(schemaResponse.schema.language.go!.name)}HandleResponse`;
       responseType = schemaResponse.schema.language.go!.responseType.name;
+      pollUntilDoneResponse = `(*${responseType}, error)`;
+      pollUntilDoneReturn = 'p.FinalResponse(ctx)';
       rawResponse = '.RawResponse';
+      unmarshalResponse = `resp.UnmarshalAsJSON(&result.${schemaResponse.schema.language.go!.responseType.value})`;
+      // for operations that do return a model add a final response method that handles the final get URL scenario
+      finalResponseDeclaration = `FinalResponse(ctx context.Context) (*${responseType}, error)`;
+      finalResponse = `FinalResponse(ctx context.Context) (*${responseType}, error) {
+        if p.pt.finalGetURL() == "" {
+          // we can end up in this situation if the async operation returns a 200
+          // with no polling URLs.  in that case return the response which should
+          // contain the JSON payload (only do this for successful terminal cases).
+          if lr := p.pt.latestResponse(); lr != nil && p.pt.hasSucceeded() {
+            result, err := p.handleResponse(lr)
+            if err != nil {
+              return nil, err
+            }
+            return result, nil
+          }
+          return nil, errors.New("missing URL for retrieving result")
+        }
+        u, err := url.Parse(p.pt.finalGetURL())
+        if err != nil {
+          return nil, err
+        }
+        req := azcore.NewRequest(http.MethodGet, *u)
+        if err != nil {
+          return nil, err
+        }
+        resp, err := p.pipeline.Do(ctx, req)
+        if err != nil {
+          return nil, err
+        }
+        return p.handleResponse(resp)
+      }`;
+      handleResponse = `
+      func (p *${pollerName}) handleResponse(resp *azcore.Response) (*${responseType}, error) {
+        result := ${responseType}{RawResponse: resp.Response}
+        if (resp.HasStatusCode(http.StatusNoContent)) {
+          return &result, nil
+        }
+        if !resp.HasStatusCode(pollingCodes[:]...) {
+          return nil, p.pt.handleError(resp)
+        }
+        return &result, ${unmarshalResponse}
+      }
+      `;
     }
     bodyText += `// ${pollerInterface} provides polling facilities until the operation completes
 type ${pollerInterface} interface {
 	Done() bool
 	Poll(ctx context.Context) (*http.Response, error)
-	FinalResponse(ctx context.Context) (*${responseType}, error)
+	${finalResponseDeclaration}
 	ResumeToken() (string, error)
 }
 
-type ${camelCase(schemaResponse.schema.language.go!.name)}HandleResponse func(*azcore.Response) (*${responseType}, error)
-
 type ${pollerName} struct {
 	// the client for making the request
 	pipeline azcore.Pipeline
 	// polling tracker
   pt pollingTracker
-  // use the response handler to check for accepted status codes
-	response ${responseHandler}
 }
 
 // Done returns true if there was an error or polling has reached a terminal state
@@ -78,34 +122,7 @@ func (p *${pollerName}) Poll(ctx context.Context) (*http.Response, error) {
   return nil, p.pt.pollingError()
 }
 
-func (p *${pollerName}) FinalResponse(ctx context.Context) (*${responseType}, error) {
-  if p.pt.finalGetURL() == "" {
-    // we can end up in this situation if the async operation returns a 200
-    // with no polling URLs.  in that case return the response which should
-    // contain the JSON payload (only do this for successful terminal cases).
-    if lr := p.pt.latestResponse(); lr != nil && p.pt.hasSucceeded() {
-      result, err := p.response(lr)
-      if err != nil {
-        return nil, err
-      }
-      return result, nil
-    }
-    return nil, errors.New("missing URL for retrieving result")
-  }
-  u, err := url.Parse(p.pt.finalGetURL())
-  if err != nil {
-    return nil, err
-  }
-  req := azcore.NewRequest(http.MethodGet, *u)
-  if err != nil {
-    return nil, err
-  }
-  resp, err := p.pipeline.Do(ctx, req)
-	if err != nil {
-		return nil, err
-	}
-	return p.response(resp)
-}
+func (p *${pollerName}) ${finalResponse}
 
 // ResumeToken generates the string token that can be used with the Resume${pollerInterface} method
 // on the client to create a new poller from the data held in the current poller type
@@ -120,21 +137,24 @@ func (p *${pollerName}) ResumeToken() (string, error) {
   return string(js), nil
 }
 
-func ${pollerName}PollUntilDone(ctx context.Context, p ${pollerInterface}, frequency time.Duration) (*${responseType}, error) {
-    for !p.Done() {
+func (p *${pollerName}) pollUntilDone(ctx context.Context, frequency time.Duration) ${pollUntilDoneResponse} {
+    for {
         resp, err := p.Poll(ctx)
         if err != nil {
             return nil, err
         }
+        if p.Done() {
+          break
+        }
         if delay := azcore.RetryAfter(resp); delay > 0 {
             time.Sleep(delay)
         } else {
             time.Sleep(frequency)
         }
     }
-    return p.FinalResponse(ctx)
+    return ${pollUntilDoneReturn}
 }
-
+${handleResponse}
 `;
   }
   text += imports.text();
@@ -234,6 +254,9 @@ type pollingTracker interface {
 
   // returns the cached HTTP response after a call to pollForStatus(), can be nil
   latestResponse() *azcore.Response
+
+  // converts an *azcore.Response to an error
+  handleError(resp *azcore.Response) error
 }
 
 type methodErrorHandler func(resp *azcore.Response) error
@@ -491,7 +514,11 @@ type pollingTrackerBase struct {
 		// it's ok if we didn't find a polling header, this will be handled elsewhere
 		return nil
 	}
-	
+  
+  func (pt *pollingTrackerBase) handleError(resp *azcore.Response) error {
+    return pt.errorHandler(resp)
+  }
+
 	// DELETE
 	
 	type pollingTrackerDelete struct {

diff --git a/src/transform/transform.ts b/src/transform/transform.ts
@@ -410,7 +410,7 @@ function createResponseType(codeModel: CodeModel, group: OperationGroup, op: Ope
       const description = `${name} contains the HTTP response from the call to the service endpoint`;
       const object = new ObjectSchema(name, description);
       object.language.go = object.language.default;
-      const pollUntilDone = newProperty('PollUntilDone', 'PollUntilDone will poll the service endpoint until a terminal state is reached or an error is received', newObject('func(ctx context.Context, frequency time.Duration) (*HTTPResponse, error)', 'TODO'));
+      const pollUntilDone = newProperty('PollUntilDone', 'PollUntilDone will poll the service endpoint until a terminal state is reached or an error is received', newObject('func(ctx context.Context, frequency time.Duration) (*http.Response, error)', 'TODO'));
       const getPoller = newProperty('Poller', 'Poller contains an initialized poller', newObject('HTTPPoller', 'TODO'));
       pollUntilDone.schema.language.go!.lroPointerException = true;
       getPoller.schema.language.go!.lroPointerException = true;

diff --git a/test/autorest/generated/lrogroup/lroretrys.go b/test/autorest/generated/lrogroup/lroretrys.go