feat(azure): improve mocking support for Task Queue and Coordinator

aranasoft · Apr 3, 2024 · 890b845 · 890b845
1 parent 8174d06
commit 890b845
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/azure/src/azure/ServiceBus/TaskQueue/TaskCoordinator.cs b/src/azure/src/azure/ServiceBus/TaskQueue/TaskCoordinator.cs
@@ -14,6 +14,14 @@ public class TaskCoordinator : ITaskCoordinator {
         private readonly ILogger<TaskCoordinator> _log;
         private readonly ITaskHandlerResolver _taskHandlerResolver;
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="TaskRequestQueue"/> class.
+        /// This constructor is intended for mocking purposes.
+        /// </summary>
+        protected TaskCoordinator()
+        {
+        }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="TaskCoordinator"/> class.
         /// </summary>
@@ -25,7 +33,7 @@ public TaskCoordinator(ITaskHandlerResolver taskHandlerResolver, ILogger<TaskCoo
         }
 
         /// <inheritdoc />
-        public Task ProcessQueueMessageAsync(string message, CancellationToken cancellationToken = default) {
+        public virtual Task ProcessQueueMessageAsync(string message, CancellationToken cancellationToken = default) {
             var taskRequest = DeserializeTaskRequest(message);
             return ProcessTaskAsync(taskRequest, cancellationToken);
         }
@@ -36,7 +44,7 @@ public Task ProcessQueueMessageAsync(string message, CancellationToken cancellat
         /// <param name="taskRequest">The <see cref="TaskRequest"/> to process.</param>
         /// <param name="cancellationToken">A <see cref="CancellationToken"/> that can be used to cancel the operation.</param>
         /// <returns>A <see cref="Task"/> that represents the asynchronous operation.</returns>
-        public async Task ProcessTaskAsync(TaskRequest taskRequest, CancellationToken cancellationToken = default) {
+        public virtual async Task ProcessTaskAsync(TaskRequest taskRequest, CancellationToken cancellationToken = default) {
             var taskRequestType = taskRequest.GetType();
             var taskRequestTrackingId = taskRequest.TrackingId.ToString("D");
             _log.LogInformation("Received task {TaskRequestTrackingId} of type {TaskRequestType}",

diff --git a/src/azure/src/azure/ServiceBus/TaskQueue/TaskRequestQueue.cs b/src/azure/src/azure/ServiceBus/TaskQueue/TaskRequestQueue.cs
@@ -37,14 +37,14 @@ public TaskRequestQueue(ServiceBusClient queueClient, string taskQueueName) {
         }
 
         /// <inheritdoc />
-        public Task AddTaskAsync(TaskRequest request, CancellationToken cancellationToken = default) {
+        public virtual Task AddTaskAsync(TaskRequest request, CancellationToken cancellationToken = default) {
             var serializedRequest = request.ToJson();
             var message = new ServiceBusMessage(Encoding.UTF8.GetBytes(serializedRequest));
             return GetQueueSender().SendMessageAsync(message, cancellationToken);
         }
 
         /// <inheritdoc />
-        public Task ScheduleTaskAsync(TaskRequest request, DateTimeOffset scheduledMessageEnqueueTime, CancellationToken cancellationToken = default) {
+        public virtual Task ScheduleTaskAsync(TaskRequest request, DateTimeOffset scheduledMessageEnqueueTime, CancellationToken cancellationToken = default) {
             var serializedRequest = request.ToJson();
             var message = new ServiceBusMessage(Encoding.UTF8.GetBytes(serializedRequest));
             return GetQueueSender().ScheduleMessageAsync(message, scheduledMessageEnqueueTime, cancellationToken);