Fix null reference error if there's no built in retry support (#1476)

* changes * addressing comments * adding null coalesce * addressing comment * nit
Azure · Apr 21, 2023 · 587714e · 587714e
1 parent ffc7914
commit 587714e
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 1 deletion.
diff --git a/src/DotNetWorker.Grpc/GrpcFunctionInvocation.cs b/src/DotNetWorker.Grpc/GrpcFunctionInvocation.cs
@@ -22,6 +22,6 @@ public GrpcFunctionInvocation(InvocationRequest invocationRequest)
 
         public override TraceContext TraceContext { get; }
 
-        public RetryContext Context => _retryContext ??= new GrpcRetryContext(_invocationRequest.RetryContext);
+        public RetryContext Context => _retryContext ??= _invocationRequest.RetryContext == null ? null! : new GrpcRetryContext(_invocationRequest.RetryContext);
     }
 }
diff --git a/test/DotNetWorkerTests/GrpcWorkerTests.cs b/test/DotNetWorkerTests/GrpcWorkerTests.cs
@@ -302,6 +302,22 @@ public async Task Invoke_SetsRetryContext()
             Assert.Equal(request.RetryContext.MaxRetryCount, _context.RetryContext.MaxRetryCount);
         }
 
+        [Fact]
+        public async Task SetRetryContextToNull()
+        {
+            var request = TestUtility.CreateInvocationRequestWithNullRetryContext();
+
+            var invocationHandler = new InvocationHandler(_mockApplication.Object,
+                _mockFeaturesFactory.Object, new JsonObjectSerializer(), _mockOutputBindingsInfoProvider.Object,
+                _mockInputConversionFeatureProvider.Object, _testLogger);
+
+            var response = await invocationHandler.InvokeAsync(request);
+
+            Assert.Equal(StatusResult.Types.Status.Success, response.Result.Status);
+            Assert.True(_context.IsDisposed);
+            Assert.Null(_context.RetryContext);
+        }
+
         [Fact]
         public async Task Invoke_CreateContextThrows_ReturnsFailure()
         {

diff --git a/test/DotNetWorkerTests/TestUtility.cs b/test/DotNetWorkerTests/TestUtility.cs
@@ -73,5 +73,19 @@ public static InvocationRequest CreateInvocationRequest(string invocationId = ""
                 }
             };
         }
+
+        public static InvocationRequest CreateInvocationRequestWithNullRetryContext(string invocationId = "")
+        {
+            return new InvocationRequest
+            {
+                InvocationId = invocationId,
+                TraceContext = new RpcTraceContext
+                {
+                    TraceParent = Guid.NewGuid().ToString(),
+                    TraceState = Guid.NewGuid().ToString()
+                },
+                RetryContext = null
+            };
+        }
     }
 }