Adding check for pollerType when resuming a poller (#384)

* Adding check for pollerType when resuming a poller * Adding test for resuming from the wrong poller Co-authored-by: Catalina Peralta <[email protected]>
Azure · May 6, 2020 · 5ec7293 · 5ec7293
1 parent dc8fb50
commit 5ec7293
Show file tree

Hide file tree

Showing 8 changed files with 195 additions and 167 deletions.
diff --git a/src/generator/operations.ts b/src/generator/operations.ts
@@ -284,7 +284,7 @@ function generateOperation(clientName: string, op: Operation, imports: ImportMan
     text += `\tif err != nil {\n`;
     text += `\t\treturn nil, err\n`;
     text += `\t}\n`;
-    text += `\tpt, err := createPollingTracker(resp, client.${info.protocolNaming.errorMethod})\n`;
+    text += `\tpt, err := createPollingTracker("${op.language.go!.pollerType.name}", resp, client.${info.protocolNaming.errorMethod})\n`;
     text += `\tif err != nil {\n`;
     text += `\t\treturn nil, err\n`;
     text += `\t}\n`;
@@ -919,7 +919,7 @@ function addResumePollerMethod(op: Operation, clientName: string): string {
   const pollerName = pascalCase(op.language.go!.pollerType.name);
   const info = <OperationNaming>op.language.go!;
   let text = `func (client *${clientName}) Resume${pollerName}(token string) (${pollerName}, error) {\n`;
-  text += `\tpt, err := resumePollingTracker(token, client.${info.protocolNaming.errorMethod})\n`;
+  text += `\tpt, err := resumePollingTracker("${op.language.go!.pollerType.name}", token, client.${info.protocolNaming.errorMethod})\n`;
   text += `\tif err != nil {\n`;
   text += `\t\treturn nil, err\n`;
   text += `\t}\n`;

diff --git a/src/generator/pollers.ts b/src/generator/pollers.ts
@@ -234,7 +234,10 @@ export async function generatePollersHelper(session: Session<CodeModel>): Promis
 
 	type pollingTrackerBase struct {
 		// resp is the last response, either from the submission of the LRO or from polling
-		resp *azcore.Response
+    resp *azcore.Response
+    
+    // PollerType is the name of the type of poller that is created
+		PollerType string \`json:"pollerType"\`
 
 	// errorHandler is the method to invoke to unmarshall an error response
 	errorHandler methodErrorHandler
@@ -698,17 +701,17 @@ export async function generatePollersHelper(session: Session<CodeModel>): Promis
 	}
 	
 	// creates a polling tracker based on the verb of the original request
-	func createPollingTracker(resp *azcore.Response, errorHandler methodErrorHandler) (pollingTracker, error) {
+	func createPollingTracker(pollerType string, resp *azcore.Response, errorHandler methodErrorHandler) (pollingTracker, error) {
 		var pt pollingTracker
 		switch strings.ToUpper(resp.Request.Method) {
 		case http.MethodDelete:
-			pt = &pollingTrackerDelete{pollingTrackerBase: pollingTrackerBase{resp: resp, errorHandler: errorHandler}}
+			pt = &pollingTrackerDelete{pollingTrackerBase: pollingTrackerBase{PollerType: pollerType, resp: resp, errorHandler: errorHandler}}
 		case http.MethodPatch:
-			pt = &pollingTrackerPatch{pollingTrackerBase: pollingTrackerBase{resp: resp, errorHandler: errorHandler}}
+			pt = &pollingTrackerPatch{pollingTrackerBase: pollingTrackerBase{PollerType: pollerType, resp: resp, errorHandler: errorHandler}}
 		case http.MethodPost:
-			pt = &pollingTrackerPost{pollingTrackerBase: pollingTrackerBase{resp: resp, errorHandler: errorHandler}}
+			pt = &pollingTrackerPost{pollingTrackerBase: pollingTrackerBase{PollerType: pollerType, resp: resp, errorHandler: errorHandler}}
 		case http.MethodPut:
-			pt = &pollingTrackerPut{pollingTrackerBase: pollingTrackerBase{resp: resp, errorHandler: errorHandler}}
+			pt = &pollingTrackerPut{pollingTrackerBase: pollingTrackerBase{PollerType: pollerType, resp: resp, errorHandler: errorHandler}}
 		default:
 			return nil, fmt.Errorf("unsupported HTTP method %s", resp.Request.Method)
 		}
@@ -722,12 +725,15 @@ export async function generatePollersHelper(session: Session<CodeModel>): Promis
 	}
 
 // creates a polling tracker from a resume token
-func resumePollingTracker(token string, errorHandler methodErrorHandler) (pollingTracker, error) {
+func resumePollingTracker(pollerType string, token string, errorHandler methodErrorHandler) (pollingTracker, error) {
 	// unmarshal into JSON object to determine the tracker type
 	obj := map[string]interface{}{}
 	err := json.Unmarshal([]byte(token), &obj)
 	if err != nil {
 		return nil, err
+  }
+  if obj["pollerType"] != pollerType {
+		return nil, fmt.Errorf("cannot resume from this poller type")
 	}
 	if obj["method"] == nil {
 		return nil, fmt.Errorf("token is missing 'method' property")

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