From 1ec8e681a44b9e18e44edf90e4b345f5613acf39 Mon Sep 17 00:00:00 2001
From: James Newton-King <james@newtonking.com>
Date: Fri, 13 Oct 2023 20:46:20 +0800
Subject: [PATCH] Fix tests

---
 .../Internal/GrpcProtocolHelpers.cs               | 15 ++++++++++++++-
 .../Internal/HttpContextStreamReader.cs           |  2 +-
 .../Internal/HttpContextStreamWriter.cs           |  2 +-
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/Grpc.AspNetCore.Server/Internal/GrpcProtocolHelpers.cs b/src/Grpc.AspNetCore.Server/Internal/GrpcProtocolHelpers.cs
index 11aa0d1a7..ef5a7b191 100644
--- a/src/Grpc.AspNetCore.Server/Internal/GrpcProtocolHelpers.cs
+++ b/src/Grpc.AspNetCore.Server/Internal/GrpcProtocolHelpers.cs
@@ -1,4 +1,4 @@
-#region Copyright notice and license
+#region Copyright notice and license
 
 // Copyright 2019 The gRPC Authors
 //
@@ -234,4 +234,17 @@ internal static bool ShouldSkipHeader(string name)
     {
         return name.StartsWith(':') || GrpcProtocolConstants.FilteredHeaders.Contains(name);
     }
+
+    internal static IHttpRequestLifetimeFeature GetRequestLifetimeFeature(HttpContext httpContext)
+    {
+        var lifetimeFeature = httpContext.Features.Get<IHttpRequestLifetimeFeature>();
+        if (lifetimeFeature is null)
+        {
+            // This should only run in tests where the HttpContext is manually created.
+            lifetimeFeature = new HttpRequestLifetimeFeature();
+            httpContext.Features.Set(lifetimeFeature);
+        }
+
+        return lifetimeFeature;
+    }
 }
diff --git a/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamReader.cs b/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamReader.cs
index b893c4c19..528aeed87 100644
--- a/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamReader.cs
+++ b/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamReader.cs
@@ -45,7 +45,7 @@ public HttpContextStreamReader(HttpContextServerCallContext serverCallContext, F
         // This is done to avoid a race condition when reading them from HttpContext later when running in a separate thread.
         _bodyReader = _serverCallContext.HttpContext.Request.BodyReader;
         // Copy lifetime feature because HttpContext.RequestAborted on .NET 6 doesn't return the real cancellation token.
-        _requestLifetimeFeature = _serverCallContext.HttpContext.Features.Get<IHttpRequestLifetimeFeature>()!;
+        _requestLifetimeFeature = GrpcProtocolHelpers.GetRequestLifetimeFeature(_serverCallContext.HttpContext);
     }
 
     public TRequest Current { get; private set; } = default!;
diff --git a/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamWriter.cs b/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamWriter.cs
index 2c8b3d286..b6a67f9fd 100644
--- a/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamWriter.cs
+++ b/src/Grpc.AspNetCore.Server/Internal/HttpContextStreamWriter.cs
@@ -48,7 +48,7 @@ public HttpContextStreamWriter(HttpContextServerCallContext context, Action<TRes
         // This is done to avoid a race condition when reading them from HttpContext later when running in a separate thread.
         _bodyWriter = context.HttpContext.Response.BodyWriter;
         // Copy lifetime feature because HttpContext.RequestAborted on .NET 6 doesn't return the real cancellation token.
-        _requestLifetimeFeature = context.HttpContext.Features.Get<IHttpRequestLifetimeFeature>()!;
+        _requestLifetimeFeature = GrpcProtocolHelpers.GetRequestLifetimeFeature(context.HttpContext);
     }
 
     public WriteOptions? WriteOptions