[FR] Add post response to Operation classes (Azure#24561)

maririos · web-flow · commit 574cca2a158b · 2021-10-12T15:28:36.000-07:00
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/src/AnalyzeDocumentOperation.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/src/AnalyzeDocumentOperation.cs
@@ -103,11 +103,12 @@ public AnalyzeDocumentOperation(string operationId, DocumentAnalysisClient clien
         /// <param name="serviceClient">The client for communicating with the Form Recognizer Azure Cognitive Service through its REST API.</param>
         /// <param name="diagnostics">The client diagnostics for exception creation in case of failure.</param>
         /// <param name="operationLocation">The address of the long-running operation. It can be obtained from the response headers upon starting the operation.</param>
-        internal AnalyzeDocumentOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation)
+        /// <param name="postResponse">Response from the POSt request that initiated the operation.</param>
+        internal AnalyzeDocumentOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation, Response postResponse)
         {
             _serviceClient = serviceClient;
             _diagnostics = diagnostics;
-            _operationInternal = new(_diagnostics, this, rawResponse: null);
+            _operationInternal = new(_diagnostics, this, rawResponse: postResponse);
 
             // TODO: Use regex to parse ids.
             // https://github.com/Azure/azure-sdk-for-net/issues/11505
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/src/BuildModelOperation.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/src/BuildModelOperation.cs
@@ -94,12 +94,13 @@ public override async ValueTask<Response<DocumentModel>> WaitForCompletionAsync(
 
         internal BuildModelOperation(
             string location,
+            Response postResponse,
             DocumentAnalysisRestClient allOperations,
             ClientDiagnostics diagnostics)
         {
             _serviceClient = allOperations;
             _diagnostics = diagnostics;
-            _operationInternal = new(_diagnostics, this, rawResponse: null, nameof(BuildModelOperation));
+            _operationInternal = new(_diagnostics, this, rawResponse: postResponse);
 
             Id = location.Split('/').Last().Split('?').FirstOrDefault();
         }
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/src/CopyModelOperation.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/src/CopyModelOperation.cs
@@ -79,11 +79,12 @@ public CopyModelOperation(string operationId, DocumentModelAdministrationClient
         /// <param name="serviceClient">The client for communicating with the Form Recognizer Azure Cognitive Service through its REST API.</param>
         /// <param name="diagnostics">The client diagnostics for exception creation in case of failure.</param>
         /// <param name="operationLocation">The address of the long-running operation. It can be obtained from the response headers upon starting the operation.</param>
-        internal CopyModelOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation)
+        /// <param name="postResponse">Response from the POSt request that initiated the operation.</param>
+        internal CopyModelOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation, Response postResponse)
         {
             _serviceClient = serviceClient;
             _diagnostics = diagnostics;
-            _operationInternal = new(_diagnostics, this, rawResponse: null);
+            _operationInternal = new(_diagnostics, this, rawResponse: postResponse);
 
             Id = operationLocation.Split('/').Last().Split('?').FirstOrDefault();
         }
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/src/DocumentAnalysisClient.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/src/DocumentAnalysisClient.cs
@@ -150,7 +150,7 @@ public virtual async Task<AnalyzeDocumentOperation> StartAnalyzeDocumentAsync(st
                     document,
                     cancellationToken).ConfigureAwait(false);
 
-                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
@@ -200,7 +200,7 @@ public virtual AnalyzeDocumentOperation StartAnalyzeDocument(string modelId, Str
                     document,
                     cancellationToken);
 
-                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
@@ -250,7 +250,7 @@ public virtual async Task<AnalyzeDocumentOperation> StartAnalyzeDocumentFromUriA
                     request,
                     cancellationToken).ConfigureAwait(false);
 
-                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
@@ -300,7 +300,7 @@ public virtual AnalyzeDocumentOperation StartAnalyzeDocumentFromUri(string model
                     request,
                     cancellationToken);
 
-                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/src/DocumentModelAdministrationClient.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/src/DocumentModelAdministrationClient.cs
@@ -148,7 +148,7 @@ public virtual BuildModelOperation StartBuildModel(Uri trainingFilesUri, string
                 };
 
                 var response = ServiceClient.DocumentAnalysisBuildDocumentModel(request, cancellationToken);
-                return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);
+                return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);
             }
             catch (Exception e)
             {
@@ -194,7 +194,7 @@ public virtual async Task<BuildModelOperation> StartBuildModelAsync(Uri training
                 };
 
                 var response = await ServiceClient.DocumentAnalysisBuildDocumentModelAsync(request, cancellationToken).ConfigureAwait(false);
-                return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);
+                return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);
             }
             catch (Exception e)
             {
@@ -603,7 +603,7 @@ public virtual CopyModelOperation StartCopyModel(string modelId, CopyAuthorizati
             try
             {
                 var response = ServiceClient.DocumentAnalysisCopyDocumentModelTo(modelId, target, cancellationToken);
-                return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
@@ -632,7 +632,7 @@ public virtual async Task<CopyModelOperation> StartCopyModelAsync(string modelId
             try
             {
                 var response = await ServiceClient.DocumentAnalysisCopyDocumentModelToAsync(modelId, target, cancellationToken).ConfigureAwait(false);
-                return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);
+                return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());
             }
             catch (Exception e)
             {
@@ -739,7 +739,7 @@ public virtual BuildModelOperation StartCreateComposedModel(IEnumerable<string>
                 };
 
                 var response = ServiceClient.DocumentAnalysisComposeDocumentModel(composeRequest, cancellationToken);
-                return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);
+                return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);
             }
             catch (Exception e)
             {
@@ -778,7 +778,7 @@ public virtual async Task<BuildModelOperation> StartCreateComposedModelAsync(IEn
                 };
 
                 var response = await ServiceClient.DocumentAnalysisComposeDocumentModelAsync(composeRequest, cancellationToken).ConfigureAwait(false);
-                return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);
+                return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);
             }
             catch (Exception e)
             {
diff --git a/sdk/formrecognizer/Azure.AI.FormRecognizer/tests/DocumentAnalysisModels/OperationsLiveTests.cs b/sdk/formrecognizer/Azure.AI.FormRecognizer/tests/DocumentAnalysisModels/OperationsLiveTests.cs
@@ -38,6 +38,7 @@ public async Task AnalyzeDocumentOperationCanPollFromNewObject()
 
             var uri = DocumentAnalysisTestEnvironment.CreateUri(TestFile.Blank);
             var operation = await client.StartAnalyzeDocumentFromUriAsync(modelId, uri);
+            Assert.IsNotNull(operation.GetRawResponse());
 
             var sameOperation = InstrumentOperation(new AnalyzeDocumentOperation(operation.Id, nonInstrumentedClient));
             await sameOperation.WaitForCompletionAsync();
@@ -54,6 +55,7 @@ public async Task BuildModelOperationCanPollFromNewObject()
             var modelId = Recording.GenerateId();
 
             var operation = await client.StartBuildModelAsync(trainingFilesUri, modelId);
+            Assert.IsNotNull(operation.GetRawResponse());
 
             var sameOperation = InstrumentOperation(new BuildModelOperation(operation.Id, nonInstrumentedClient));
             await sameOperation.WaitForCompletionAsync();
@@ -90,6 +92,7 @@ public async Task CopyModelOperationCanPollFromNewObject()
             CopyAuthorization targetAuth = await client.GetCopyAuthorizationAsync(targetModelId);
 
             var operation = await client.StartCopyModelAsync(trainedModel.ModelId, targetAuth);
+            Assert.IsNotNull(operation.GetRawResponse());
 
             var sameOperation = InstrumentOperation(new CopyModelOperation(operation.Id, nonInstrumentedClient));
             await sameOperation.WaitForCompletionAsync();

Original file line number	Diff line number	Diff line change
`@@ -94,12 +94,13 @@ public override async ValueTask<Response<DocumentModel>> WaitForCompletionAsync(`
`94`	`94`
`95`	`95`	`internal BuildModelOperation(`
`96`	`96`	`string location,`
	`97`	`+ Response postResponse,`
`97`	`98`	`DocumentAnalysisRestClient allOperations,`
`98`	`99`	`ClientDiagnostics diagnostics)`
`99`	`100`	`{`
`100`	`101`	`_serviceClient = allOperations;`
`101`	`102`	`_diagnostics = diagnostics;`
`102`		`- _operationInternal = new(_diagnostics, this, rawResponse: null, nameof(BuildModelOperation));`
	`103`	`+ _operationInternal = new(_diagnostics, this, rawResponse: postResponse);`
`103`	`104`
`104`	`105`	`Id = location.Split('/').Last().Split('?').FirstOrDefault();`
`105`	`106`	`}`
Original file line number	Diff line number	Diff line change
`@@ -79,11 +79,12 @@ public CopyModelOperation(string operationId, DocumentModelAdministrationClient`
`79`	`79`	`/// <param name="serviceClient">The client for communicating with the Form Recognizer Azure Cognitive Service through its REST API.</param>`
`80`	`80`	`/// <param name="diagnostics">The client diagnostics for exception creation in case of failure.</param>`
`81`	`81`	`/// <param name="operationLocation">The address of the long-running operation. It can be obtained from the response headers upon starting the operation.</param>`
`82`		`- internal CopyModelOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation)`
	`82`	`+ /// <param name="postResponse">Response from the POSt request that initiated the operation.</param>`
	`83`	`+ internal CopyModelOperation(DocumentAnalysisRestClient serviceClient, ClientDiagnostics diagnostics, string operationLocation, Response postResponse)`
`83`	`84`	`{`
`84`	`85`	`_serviceClient = serviceClient;`
`85`	`86`	`_diagnostics = diagnostics;`
`86`		`- _operationInternal = new(_diagnostics, this, rawResponse: null);`
	`87`	`+ _operationInternal = new(_diagnostics, this, rawResponse: postResponse);`
`87`	`88`
`88`	`89`	`Id = operationLocation.Split('/').Last().Split('?').FirstOrDefault();`
`89`	`90`	`}`
Original file line number	Diff line number	Diff line change
`@@ -150,7 +150,7 @@ public virtual async Task<AnalyzeDocumentOperation> StartAnalyzeDocumentAsync(st`
`150`	`150`	`document,`
`151`	`151`	`cancellationToken).ConfigureAwait(false);`
`152`	`152`
`153`		`- return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`153`	`+ return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`154`	`154`	`}`
`155`	`155`	`catch (Exception e)`
`156`	`156`	`{`
`@@ -200,7 +200,7 @@ public virtual AnalyzeDocumentOperation StartAnalyzeDocument(string modelId, Str`
`200`	`200`	`document,`
`201`	`201`	`cancellationToken);`
`202`	`202`
`203`		`- return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`203`	`+ return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`204`	`204`	`}`
`205`	`205`	`catch (Exception e)`
`206`	`206`	`{`
`@@ -250,7 +250,7 @@ public virtual async Task<AnalyzeDocumentOperation> StartAnalyzeDocumentFromUriA`
`250`	`250`	`request,`
`251`	`251`	`cancellationToken).ConfigureAwait(false);`
`252`	`252`
`253`		`- return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`253`	`+ return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`254`	`254`	`}`
`255`	`255`	`catch (Exception e)`
`256`	`256`	`{`
`@@ -300,7 +300,7 @@ public virtual AnalyzeDocumentOperation StartAnalyzeDocumentFromUri(string model`
`300`	`300`	`request,`
`301`	`301`	`cancellationToken);`
`302`	`302`
`303`		`- return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`303`	`+ return new AnalyzeDocumentOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`304`	`304`	`}`
`305`	`305`	`catch (Exception e)`
`306`	`306`	`{`
Original file line number	Diff line number	Diff line change
`@@ -148,7 +148,7 @@ public virtual BuildModelOperation StartBuildModel(Uri trainingFilesUri, string`
`148`	`148`	`};`
`149`	`149`
`150`	`150`	`var response = ServiceClient.DocumentAnalysisBuildDocumentModel(request, cancellationToken);`
`151`		`- return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);`
	`151`	`+ return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);`
`152`	`152`	`}`
`153`	`153`	`catch (Exception e)`
`154`	`154`	`{`
`@@ -194,7 +194,7 @@ public virtual async Task<BuildModelOperation> StartBuildModelAsync(Uri training`
`194`	`194`	`};`
`195`	`195`
`196`	`196`	`var response = await ServiceClient.DocumentAnalysisBuildDocumentModelAsync(request, cancellationToken).ConfigureAwait(false);`
`197`		`- return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);`
	`197`	`+ return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);`
`198`	`198`	`}`
`199`	`199`	`catch (Exception e)`
`200`	`200`	`{`
`@@ -603,7 +603,7 @@ public virtual CopyModelOperation StartCopyModel(string modelId, CopyAuthorizati`
`603`	`603`	`try`
`604`	`604`	`{`
`605`	`605`	`var response = ServiceClient.DocumentAnalysisCopyDocumentModelTo(modelId, target, cancellationToken);`
`606`		`- return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`606`	`+ return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`607`	`607`	`}`
`608`	`608`	`catch (Exception e)`
`609`	`609`	`{`
`@@ -632,7 +632,7 @@ public virtual async Task<CopyModelOperation> StartCopyModelAsync(string modelId`
`632`	`632`	`try`
`633`	`633`	`{`
`634`	`634`	`var response = await ServiceClient.DocumentAnalysisCopyDocumentModelToAsync(modelId, target, cancellationToken).ConfigureAwait(false);`
`635`		`- return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation);`
	`635`	`+ return new CopyModelOperation(ServiceClient, Diagnostics, response.Headers.OperationLocation, response.GetRawResponse());`
`636`	`636`	`}`
`637`	`637`	`catch (Exception e)`
`638`	`638`	`{`
`@@ -739,7 +739,7 @@ public virtual BuildModelOperation StartCreateComposedModel(IEnumerable<string>`
`739`	`739`	`};`
`740`	`740`
`741`	`741`	`var response = ServiceClient.DocumentAnalysisComposeDocumentModel(composeRequest, cancellationToken);`
`742`		`- return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);`
	`742`	`+ return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);`
`743`	`743`	`}`
`744`	`744`	`catch (Exception e)`
`745`	`745`	`{`
`@@ -778,7 +778,7 @@ public virtual async Task<BuildModelOperation> StartCreateComposedModelAsync(IEn`
`778`	`778`	`};`
`779`	`779`
`780`	`780`	`var response = await ServiceClient.DocumentAnalysisComposeDocumentModelAsync(composeRequest, cancellationToken).ConfigureAwait(false);`
`781`		`- return new BuildModelOperation(response.Headers.OperationLocation, ServiceClient, Diagnostics);`
	`781`	`+ return new BuildModelOperation(response.Headers.OperationLocation, response.GetRawResponse(), ServiceClient, Diagnostics);`
`782`	`782`	`}`
`783`	`783`	`catch (Exception e)`
`784`	`784`	`{`