diff --git a/test/AspNetCoreModule.Test/FunctionalTestHelper.cs b/test/AspNetCoreModule.Test/FunctionalTestHelper.cs
index 055e850..32397eb 100644
--- a/test/AspNetCoreModule.Test/FunctionalTestHelper.cs
+++ b/test/AspNetCoreModule.Test/FunctionalTestHelper.cs
@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using AspNetCoreModule.Test.Framework;
@@ -1536,9 +1536,6 @@ namespace AspNetCoreModule.Test
                         testSite.AspNetCoreApp.CreateFile(new string[] { "test" }, "App_Offline.Htm");
                         Thread.Sleep(1000);
 
-                        // ToDo: This should be replaced when the server can handle this automaticially
-                        // send a websocket data to invoke the server side websocket disconnection after the app_offline
-                        websocketClient.SendTextData("test");
                         bool connectionClosedFromServer = websocketClient.WaitForWebSocketState(WebSocketState.ConnectionClosed);
 
                         if (connectionClosedFromServer)
@@ -1596,11 +1593,6 @@ namespace AspNetCoreModule.Test
                             Thread.Sleep(1000);
                         }
                         
-                        // ToDo: This should be replaced when the server can handle this automaticially
-                        // send a websocket data to invoke the server side websocket disconnection after the app_offline
-                        websocketClient.SendTextData("test");
-                        bool connectionClosedFromServer = websocketClient.WaitForWebSocketState(WebSocketState.ConnectionClosed);
-                        
                         // Verify server side connection closing is done successfully
                         Assert.True(connectionClosedFromServer, "Closing Handshake initiated from Server");
 
@@ -1874,10 +1866,6 @@ namespace AspNetCoreModule.Test
                                     // put app_offline
                                     testSite.AspNetCoreApp.CreateFile(new string[] { "test" }, "App_Offline.Htm");
                                     Thread.Sleep(1000);
-
-                                    // ToDo: remove this when server can handle this automatically
-                                    // send a websocket data to invoke the server side websocket disconnection after the app_offline
-                                    websocketClient.SendTextData("test");
                                     
                                     // wait for the gracefulshutdown finished
                                     Thread.Sleep(shutdownDelayTime);
diff --git a/test/AspNetCoreModule.TestSites.Standard/Program.cs b/test/AspNetCoreModule.TestSites.Standard/Program.cs
index 6c3f103..f998e20 100644
--- a/test/AspNetCoreModule.TestSites.Standard/Program.cs
+++ b/test/AspNetCoreModule.TestSites.Standard/Program.cs
@@ -15,7 +15,6 @@ namespace AspnetCoreModule.TestSites.Standard
     public static class Program
     {
         public static IApplicationLifetime AappLifetime;
-        public static bool AappLifetimeStopping = false;
         public static int GracefulShutdownDelayTime = 0; 
 
         private static X509Certificate2 _x509Certificate2;
@@ -141,7 +140,7 @@ namespace AspnetCoreModule.TestSites.Standard
             );
             AappLifetime.ApplicationStopping.Register(
                 () => {
-                    AappLifetimeStopping = true;
+                    WebSocketConnections.CloseAll();
                     Thread.Sleep(Startup.SleeptimeWhileClosing / 2);
                     Thread.Sleep(GracefulShutdownDelayTime);
                 }
diff --git a/test/AspNetCoreModule.TestSites.Standard/Startup.cs b/test/AspNetCoreModule.TestSites.Standard/Startup.cs
index bcb41c2..54fdf6b 100644
--- a/test/AspNetCoreModule.TestSites.Standard/Startup.cs
+++ b/test/AspNetCoreModule.TestSites.Standard/Startup.cs
@@ -33,11 +33,14 @@ namespace AspnetCoreModule.TestSites.Standard
 
         private async Task Echo(WebSocket webSocket)
         {
+            int webSocketIndex = WebSocketConnections.GetLastIndex();
+            WebSocketConnections.WebSockets.Add(webSocketIndex, webSocket);
+
             var buffer = new byte[1024 * 4];
             var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
             bool closeFromServer = false;
-            string closeFromServerCmd = "CloseFromServer";
-            string closingFromServer = "ClosingFromServer";
+            string closeFromServerCmd = WebSocketConnections.CloseFromServerCmd;
+            string closingFromServer = WebSocketConnections.ClosingFromServer;
             int closeFromServerLength = closeFromServerCmd.Length;
 
             bool echoBack = true;
@@ -45,11 +48,9 @@ namespace AspnetCoreModule.TestSites.Standard
 
             while (!result.CloseStatus.HasValue)
             {
-                if ((result.Count == closeFromServerLength && System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd) 
-                    || Program.AappLifetimeStopping == true)
+                if ((result.Count == closeFromServerLength && System.Text.Encoding.ASCII.GetString(buffer).Substring(0, result.Count) == closeFromServerCmd))
                 {
                     // start closing handshake from backend process when client send "CloseFromServer" text message 
-                    // or when any message is sent from client during the graceful shutdown.
                     closeFromServer = true;
                     await webSocket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, closingFromServer, CancellationToken.None);
                 }
diff --git a/test/AspNetCoreModule.TestSites.Standard/WebSocketConnections.cs b/test/AspNetCoreModule.TestSites.Standard/WebSocketConnections.cs
new file mode 100644
index 0000000..ee6d0d7
--- /dev/null
+++ b/test/AspNetCoreModule.TestSites.Standard/WebSocketConnections.cs
@@ -0,0 +1,36 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net.WebSockets;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace AspnetCoreModule.TestSites.Standard
+{
+    public class WebSocketConnections
+    {
+        private static int lastIndex = 0;
+        private static Object thisLock = new Object();
+        public static string CloseFromServerCmd = "CloseFromServer";
+        public static string ClosingFromServer = "ClosingFromServer";
+
+        public static Dictionary<int, WebSocket> WebSockets = new Dictionary<int, WebSocket>();
+
+        public static int GetLastIndex()
+        {
+            lock (thisLock)
+            {
+                lastIndex++;
+                return lastIndex;
+            }
+        }
+
+        public async static void CloseAll()
+        {
+            foreach (KeyValuePair<int, WebSocket> entry in WebSockets)
+            {
+                await entry.Value.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, ClosingFromServer, CancellationToken.None);
+            }
+        }
+    }
+}