From 736cc326bf6843881c0073a3dbd7b57a93a33e77 Mon Sep 17 00:00:00 2001
From: Luke Bakken <luke@bakken.io>
Date: Mon, 22 Apr 2024 12:09:41 -0700
Subject: [PATCH] Close connections using HTTP API

---
 projects/Test/Common/Common.csproj            |  1 +
 projects/Test/Common/RabbitMQCtl.cs           | 48 +++++++++++++++++--
 .../Test/Common/TestConnectionRecoveryBase.cs |  2 +-
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/projects/Test/Common/Common.csproj b/projects/Test/Common/Common.csproj
index 59c7b946b5..3e98ffc6d9 100644
--- a/projects/Test/Common/Common.csproj
+++ b/projects/Test/Common/Common.csproj
@@ -26,6 +26,7 @@
   </ItemGroup>
   
   <ItemGroup>
+    <PackageReference Include="EasyNetQ.Management.Client" Version="2.0.0" />
     <PackageReference Include="xunit" Version="2.7.0" />
     <PackageReference Include="xunit.abstractions" Version="2.0.3" />
   </ItemGroup>
diff --git a/projects/Test/Common/RabbitMQCtl.cs b/projects/Test/Common/RabbitMQCtl.cs
index 751a273559..b26e61add6 100644
--- a/projects/Test/Common/RabbitMQCtl.cs
+++ b/projects/Test/Common/RabbitMQCtl.cs
@@ -30,10 +30,13 @@
 //---------------------------------------------------------------------------
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
+using System.Linq;
 using System.Text.RegularExpressions;
 using System.Threading.Tasks;
+using EasyNetQ.Management.Client;
 using RabbitMQ.Client;
 using Xunit.Abstractions;
 
@@ -41,7 +44,9 @@ namespace Test
 {
     public class RabbitMQCtl
     {
-        private static readonly char[] newLine = new char[] { '\n' };
+        private static readonly TimeSpan s_closeConnectionDelay = TimeSpan.FromSeconds(2);
+        private static readonly ManagementClient s_managementClient;
+        private static readonly char[] s_newLine = new char[] { '\n' };
         // NOTE: \r?
         // https://learn.microsoft.com/en-us/dotnet/standard/base-types/regular-expression-options#multiline-mode
         private static readonly Regex s_getConnectionProperties =
@@ -49,15 +54,50 @@ public class RabbitMQCtl
 
         private readonly ITestOutputHelper _output;
 
+        static RabbitMQCtl()
+        {
+            var managementUri = new Uri("http://localhost:15672");
+            s_managementClient = new ManagementClient(managementUri, "guest", "guest");
+        }
+
         public RabbitMQCtl(ITestOutputHelper output)
         {
             _output = output;
         }
 
-        public async Task CloseConnectionAsync(IConnection conn)
+        public static async Task CloseConnectionAsync(IConnection conn)
         {
-            string pid = await GetConnectionPidAsync(conn.ClientProvidedName);
-            await CloseConnectionAsync(pid);
+            ushort tries = 10;
+            EasyNetQ.Management.Client.Model.Connection connectionToClose;
+            do
+            {
+                IReadOnlyList<EasyNetQ.Management.Client.Model.Connection> connections;
+                do
+                {
+                    await Task.Delay(s_closeConnectionDelay);
+                    connections = await s_managementClient.GetConnectionsAsync();
+                } while (connections.Count == 0);
+
+                connectionToClose = connections.Where(c0 =>
+                    string.Equals((string)c0.ClientProperties["connection_name"], conn.ClientProvidedName,
+                    StringComparison.InvariantCultureIgnoreCase)).FirstOrDefault();
+
+                if (connectionToClose == null)
+                {
+                    tries--;
+                }
+                else
+                {
+                    break;
+                }
+            } while (tries > 0);
+
+            if (tries == 0)
+            {
+                throw new InvalidOperationException($"Could not delete connection: '{conn.ClientProvidedName}'");
+            }
+
+            await s_managementClient.CloseConnectionAsync(connectionToClose);
         }
 
         public Task AddUserAsync(string username, string password)
diff --git a/projects/Test/Common/TestConnectionRecoveryBase.cs b/projects/Test/Common/TestConnectionRecoveryBase.cs
index 439da74f56..f220c2f08d 100644
--- a/projects/Test/Common/TestConnectionRecoveryBase.cs
+++ b/projects/Test/Common/TestConnectionRecoveryBase.cs
@@ -170,7 +170,7 @@ internal async Task<AutorecoveringConnection> CreateAutorecoveringConnectionWith
 
         protected Task CloseConnectionAsync(IConnection conn)
         {
-            return _rabbitMQCtl.CloseConnectionAsync(conn);
+            return RabbitMQCtl.CloseConnectionAsync(conn);
         }
 
         protected Task CloseAndWaitForRecoveryAsync()