diff --git a/src/Custom/Batch/CreateBatchOperation.Protocol.cs b/src/Custom/Batch/CreateBatchOperation.Protocol.cs index 7c9b4041c..77e021092 100644 --- a/src/Custom/Batch/CreateBatchOperation.Protocol.cs +++ b/src/Custom/Batch/CreateBatchOperation.Protocol.cs @@ -60,13 +60,7 @@ public static async Task RehydrateAsync(BatchClient client CreateBatchOperationToken token = CreateBatchOperationToken.FromToken(rehydrationToken); - ClientResult result = await client.GetBatchAsync(token.BatchId, cancellationToken.ToRequestOptions()).ConfigureAwait(false); - PipelineResponse response = result.GetRawResponse(); - - using JsonDocument doc = JsonDocument.Parse(response.Content); - string status = doc.RootElement.GetProperty("status"u8).GetString()!; - - return client.CreateCreateBatchOperation(token.BatchId, status, response); + return await RehydrateAsync(client, token.BatchId, cancellationToken).ConfigureAwait(false); } /// @@ -87,13 +81,53 @@ public static CreateBatchOperation Rehydrate(BatchClient client, ContinuationTok CreateBatchOperationToken token = CreateBatchOperationToken.FromToken(rehydrationToken); - ClientResult result = client.GetBatch(token.BatchId, cancellationToken.ToRequestOptions()); + return Rehydrate(client, token.BatchId, cancellationToken); + } + + /// + /// Recreates a from a batch ID. + /// + /// The used to obtain the + /// operation status from the service. + /// The ID of the batch operation to rehydrate. + /// A token that can be used to cancel the + /// request. + /// The rehydrated operation. + /// is null. + public static async Task RehydrateAsync(BatchClient client, string batchId, CancellationToken cancellationToken = default) + { + Argument.AssertNotNull(client, nameof(client)); + + ClientResult result = await client.GetBatchAsync(batchId, cancellationToken.ToRequestOptions()).ConfigureAwait(false); + PipelineResponse response = result.GetRawResponse(); + + using JsonDocument doc = JsonDocument.Parse(response.Content); + string status = doc.RootElement.GetProperty("status"u8).GetString()!; + + return client.CreateCreateBatchOperation(batchId, status, response); + } + + /// + /// Recreates a from a batch ID. + /// + /// The used to obtain the + /// operation status from the service. + /// The ID of the batch operation to rehydrate. + /// A token that can be used to cancel the + /// request. + /// The rehydrated operation. + /// is null. + public static CreateBatchOperation Rehydrate(BatchClient client, string batchId, CancellationToken cancellationToken = default) + { + Argument.AssertNotNull(client, nameof(client)); + + ClientResult result = client.GetBatch(batchId, cancellationToken.ToRequestOptions()); PipelineResponse response = result.GetRawResponse(); using JsonDocument doc = JsonDocument.Parse(response.Content); string status = doc.RootElement.GetProperty("status"u8).GetString()!; - return client.CreateCreateBatchOperation(token.BatchId, status, response); + return client.CreateCreateBatchOperation(batchId, status, response); } /// diff --git a/tests/Batch/BatchTests.cs b/tests/Batch/BatchTests.cs index f24b5477f..ea31f05f9 100644 --- a/tests/Batch/BatchTests.cs +++ b/tests/Batch/BatchTests.cs @@ -153,8 +153,9 @@ public async Task CreateGetAndCancelBatchProtocol() Assert.That(status, Is.EqualTo("validating")); } - [Test] - public async Task CanRehydrateBatchOperation() + [TestCase(true)] + [TestCase(false)] + public async Task CanRehydrateBatchOperation(bool useBatchId) { using MemoryStream testFileStream = new(); using StreamWriter streamWriter = new(testFileStream); @@ -183,13 +184,22 @@ public async Task CanRehydrateBatchOperation() ? await client.CreateBatchAsync(content, waitUntilCompleted: false) : client.CreateBatch(content, waitUntilCompleted: false); - // Simulate rehydration of the operation - BinaryData rehydrationBytes = batchOperation.RehydrationToken.ToBytes(); - ContinuationToken rehydrationToken = ContinuationToken.FromBytes(rehydrationBytes); - - CreateBatchOperation rehydratedOperation = IsAsync ? - await CreateBatchOperation.RehydrateAsync(client, rehydrationToken) : - CreateBatchOperation.Rehydrate(client, rehydrationToken); + CreateBatchOperation rehydratedOperation; + if (useBatchId) + { + rehydratedOperation = IsAsync ? + await CreateBatchOperation.RehydrateAsync(client, batchOperation.BatchId) : + CreateBatchOperation.Rehydrate(client, batchOperation.BatchId); + } + else { + // Simulate rehydration of the operation + BinaryData rehydrationBytes = batchOperation.RehydrationToken.ToBytes(); + ContinuationToken rehydrationToken = ContinuationToken.FromBytes(rehydrationBytes); + + rehydratedOperation = IsAsync ? + await CreateBatchOperation.RehydrateAsync(client, rehydrationToken) : + CreateBatchOperation.Rehydrate(client, rehydrationToken); + } static bool Validate(CreateBatchOperation operation) {