diff --git a/src/Microsoft.AspNet.TestHost/ClientHandler.cs b/src/Microsoft.AspNet.TestHost/ClientHandler.cs index a1e281a7ca..8d8855d81e 100644 --- a/src/Microsoft.AspNet.TestHost/ClientHandler.cs +++ b/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)); } diff --git a/src/Microsoft.AspNet.TestHost/ResponseFeature.cs b/src/Microsoft.AspNet.TestHost/ResponseFeature.cs index de9765b383..b8bf81dcba 100644 --- a/src/Microsoft.AspNet.TestHost/ResponseFeature.cs +++ b/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 callback, object state) + { + var prior = _responseCompleted; + _responseCompleted = () => + { + callback(state); + prior(); + }; + } + public void FireOnSendingHeaders() { _sendingHeaders(); HeadersSent = true; } + + public void FireOnResponseCompleted() + { + _responseCompleted(); + } } }