From d7056f151593fd3d1849be4302a98a322d110573 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Sun, 4 Aug 2024 17:44:12 -0300
Subject: [PATCH] tests that caught the previous two bugs.

---
 listener_test.go | 97 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 97 insertions(+)

diff --git a/listener_test.go b/listener_test.go
index a987de8..6507d51 100644
--- a/listener_test.go
+++ b/listener_test.go
@@ -1,6 +1,7 @@
 package khatru
 
 import (
+	"math/rand"
 	"testing"
 
 	"github.com/nbd-wtf/go-nostr"
@@ -404,3 +405,99 @@ func TestListenerMoreStuffWithMultipleRelays(t *testing.T) {
 		}, rlz.listeners)
 	})
 }
+
+func TestRandomListenerClientRemoving(t *testing.T) {
+	rl := NewRelay()
+
+	f := nostr.Filter{Kinds: []int{1}}
+	cancel := func(cause error) {}
+
+	websockets := make([]*WebSocket, 0, 20)
+
+	l := 0
+
+	for i := 0; i < 20; i++ {
+		ws := &WebSocket{}
+		websockets = append(websockets, ws)
+		rl.clients[ws] = nil
+	}
+
+	for j := 0; j < 20; j++ {
+		for i := 0; i < 20; i++ {
+			ws := websockets[i]
+			w := string(rune(i + 65))
+
+			if rand.Intn(2) < 1 {
+				l++
+				rl.addListener(ws, w+":"+string(rune(j+97)), rl, f, cancel)
+			}
+		}
+	}
+
+	require.Len(t, rl.clients, 20)
+	require.Len(t, rl.listeners, l)
+
+	for ws := range rl.clients {
+		rl.removeClientAndListeners(ws)
+	}
+
+	require.Len(t, rl.clients, 0)
+	require.Len(t, rl.listeners, 0)
+}
+
+func TestRandomListenerIdRemoving(t *testing.T) {
+	rl := NewRelay()
+
+	f := nostr.Filter{Kinds: []int{1}}
+	cancel := func(cause error) {}
+
+	websockets := make([]*WebSocket, 0, 20)
+
+	type wsid struct {
+		ws *WebSocket
+		id string
+	}
+
+	subs := make([]wsid, 0, 20*20)
+	extra := 0
+
+	for i := 0; i < 20; i++ {
+		ws := &WebSocket{}
+		websockets = append(websockets, ws)
+		rl.clients[ws] = nil
+	}
+
+	for j := 0; j < 20; j++ {
+		for i := 0; i < 20; i++ {
+			ws := websockets[i]
+			w := string(rune(i + 65))
+
+			if rand.Intn(2) < 1 {
+				id := w + ":" + string(rune(j+97))
+				rl.addListener(ws, id, rl, f, cancel)
+				subs = append(subs, wsid{ws, id})
+
+				if rand.Intn(5) < 1 {
+					rl.addListener(ws, id, rl, f, cancel)
+					extra++
+				}
+			}
+		}
+	}
+
+	require.Len(t, rl.clients, 20)
+	require.Len(t, rl.listeners, len(subs)+extra)
+
+	rand.Shuffle(len(subs), func(i, j int) {
+		subs[i], subs[j] = subs[j], subs[i]
+	})
+	for _, wsidToRemove := range subs {
+		rl.removeListenerId(wsidToRemove.ws, wsidToRemove.id)
+	}
+
+	require.Len(t, rl.listeners, 0)
+	require.Len(t, rl.clients, 20)
+	for _, specs := range rl.clients {
+		require.Len(t, specs, 0)
+	}
+}