From 8402cbfa8d48b0286b027c2583290b6caae93a38 Mon Sep 17 00:00:00 2001
From: Jay <wisehodl@pm.me>
Date: Wed, 15 Apr 2026 14:28:40 -0400
Subject: [PATCH] Refactored, simplified tests.

---
 pool_test.go | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/pool_test.go b/pool_test.go
index 4583dc6..9893da2 100644
--- a/pool_test.go
+++ b/pool_test.go
@@ -110,19 +110,13 @@ func TestPoolRemove(t *testing.T) {
 		pool.dialer = mockDialer
 
 		pool.Add("wss://test")
+		expectEvent(t, pool.events, "wss://test", EventConnected)
 
 		err = pool.Remove("wss://test/")
 		assert.NoError(t, err)
 
 		// expect a disconnected event
-		assert.Eventually(t, func() bool {
-			select {
-			case e := <-pool.events:
-				return e.URL == "wss://test" && e.Kind == EventDisconnected
-			default:
-				return false
-			}
-		}, testTimeout, testTick, "expected disconnected event")
+		expectEvent(t, pool.events, "wss://test", EventDisconnected)
 
 		// connection no longer in pool
 		pool.mu.Lock()
@@ -143,9 +137,6 @@ func TestPoolRemove(t *testing.T) {
 		assert.NoError(t, err)
 		pool.dialer = mockDialer
 
-		pool.Add("wss://peer1")
-		pool.Add("wss://peer2")
-
 		// remove unknown connection
 		err = pool.Remove("wss://unknown")
 		assert.ErrorContains(t, err, "connection not found")
@@ -163,15 +154,31 @@ func TestPoolRemove(t *testing.T) {
 		assert.NoError(t, err)
 		pool.dialer = mockDialer
 
-		pool.Add("wss://peer1")
-		pool.Add("wss://peer2")
-
 		// close pool
 		pool.Close()
 
 		// attempt to remove connection
-		err = pool.Remove("wss://peer1")
+		err = pool.Remove("wss://test")
 		assert.ErrorContains(t, err, "pool is closed")
 	})
 
 }
+
+func expectEvent(
+	t *testing.T,
+	events chan PoolEvent,
+	expectedURL string,
+	expectedKind PoolEventKind,
+) {
+	t.Helper()
+	assert.Eventually(t, func() bool {
+		select {
+		case e := <-events:
+			return e.URL == expectedURL && e.Kind == expectedKind
+		default:
+			return false
+		}
+	}, testTimeout, testTick,
+		fmt.Sprintf("expected event: URL=%q, Kind=%q",
+			expectedURL, expectedKind.String()))
+}