From 71c58da4000558fa268b79276123fdfb916f77f0 Mon Sep 17 00:00:00 2001
From: James Newton-King <james@newtonking.com>
Date: Tue, 1 Oct 2024 10:42:10 +0800
Subject: [PATCH] Cleanup gRPC unit testing helpers in tester sample

---
 .../UnitTests/Helpers/TestServerCallContext.cs   |  5 +----
 .../UnitTests/Helpers/TestServerStreamWriter.cs  | 16 +++++++++++-----
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerCallContext.cs b/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerCallContext.cs
index 1f4cbd3bf..ddc4aa405 100644
--- a/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerCallContext.cs
+++ b/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerCallContext.cs
@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -22,8 +22,6 @@
 using System.Threading.Tasks;
 using Grpc.Core;
 
-// TODO(JamesNK): Remove nullable override after Grpc.Core.Api update
-#pragma warning disable CS8764 // Nullability of return type doesn't match overridden member (possibly because of nullability attributes).
 namespace Tests.Server.UnitTests.Helpers
 {
     public class TestServerCallContext : ServerCallContext
@@ -81,4 +79,3 @@ public static TestServerCallContext Create(Metadata? requestHeaders = null, Canc
         }
     }
 }
-#pragma warning restore CS8764 // Nullability of return type doesn't match overridden member (possibly because of nullability attributes).
diff --git a/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerStreamWriter.cs b/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerStreamWriter.cs
index 08163f12a..a00d8cf79 100644
--- a/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerStreamWriter.cs
+++ b/examples/Tester/Tests/Server/UnitTests/Helpers/TestServerStreamWriter.cs
@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -22,8 +22,6 @@
 using System.Threading.Tasks;
 using Grpc.Core;
 
-// TODO(JamesNK): Remove nullable override after Grpc.Core.Api update
-#pragma warning disable CS8766 // Nullability of reference types in return type doesn't match implicitly implemented member (possibly because of nullability attributes).
 namespace Tests.Server.UnitTests.Helpers
 {
     public class TestServerStreamWriter<T> : IServerStreamWriter<T> where T : class
@@ -63,8 +61,12 @@ public IAsyncEnumerable<T> ReadAllAsync()
             }
         }
 
-        public Task WriteAsync(T message)
+        public Task WriteAsync(T message, CancellationToken cancellationToken)
         {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return Task.FromCanceled(cancellationToken);
+            }
             if (_serverCallContext.CancellationToken.IsCancellationRequested)
             {
                 return Task.FromCanceled(_serverCallContext.CancellationToken);
@@ -77,6 +79,10 @@ public Task WriteAsync(T message)
 
             return Task.CompletedTask;
         }
+
+        public Task WriteAsync(T message)
+        {
+            return WriteAsync(message, CancellationToken.None);
+        }
     }
 }
-#pragma warning restore CS8766 // Nullability of reference types in return type doesn't match implicitly implemented member (possibly because of nullability attributes).