React to aspnet/HttpAbstractions#160 - Implementing OnResponseCompleted in TestHost

src/Microsoft.AspNet.TestHost/ClientHandler.cs

@@ -174,6 +174,7 @@ namespace Microsoft.AspNet.TestHost
                 if (!_responseTcs.Task.IsCompleted)
                 {
                     var response = GenerateResponse();
+                    _responseFeature.FireOnResponseCompleted();
                     // Dispatch, as TrySetResult will synchronously execute the waiters callback and block our Write.
                     Task.Factory.StartNew(() => _responseTcs.TrySetResult(response));
                 }

src/Microsoft.AspNet.TestHost/ResponseFeature.cs

@@ -11,6 +11,7 @@ namespace Microsoft.AspNet.TestHost
     internal class ResponseFeature : IHttpResponseFeature
     {
         private Action _sendingHeaders = () => { };
+        private Action _responseCompleted = () => { };
 
         public ResponseFeature()
         {
@@ -42,10 +43,25 @@ namespace Microsoft.AspNet.TestHost
             };
         }
 
+        public void OnResponseCompleted(Action<object> callback, object state)
+        {
+            var prior = _responseCompleted;
+            _responseCompleted = () =>
+            {
+                callback(state);
+                prior();
+            };
+        }
+
         public void FireOnSendingHeaders()
         {
             _sendingHeaders();
             HeadersSent = true;
         }
+
+        public void FireOnResponseCompleted()
+        {
+            _responseCompleted();
+        }
     }
 }