Adds tests for closing channel and connection registry

channel/channel_test.go

@@ -1,6 +1,7 @@
 package channel
 
 import (
+	"context"
 	"net/http"
 	"testing"
 
@@ -132,3 +133,16 @@ func TestChannel_WithOriginPatterns(t *testing.T) {
 		t.Errorf("Unexpected to get default origin pattern: got %s, expected %s", channel.config.originPatterns[1], "test2")
 	}
 }
+func TestChannel_Shutdown(t *testing.T) {
+	path := "/test/path"
+	dispatcher := mocks.NewMockDispatcher(t)
+
+	channel := NewChannel(path, dispatcher, NewConnectionRegistry())
+
+	// Call the Shutdown method
+	err := channel.Shutdown(context.Background())
+
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}

channel/connection_registry.go

@@ -107,8 +107,8 @@ func (r *ConnectionRegistry) Shutdown(ctx context.Context) error {
 		wg.Add(1)
 
 		go func() {
+			defer wg.Done()
 			c.Close(ctx, websocket.StatusServiceRestart, "")
-			wg.Done()
 		}()
 	}
 

channel/connection_registry_test.go

@@ -95,3 +95,32 @@ func TestConnectionRegistry_WithMaxFrameLimit(t *testing.T) {
 		t.Errorf("Unexpected frame size limit: got %d, expected %d", registry.frameSizeLimit, 100)
 	}
 }
+func TestConnectionRegistry_Shutdown(t *testing.T) {
+	ctx := context.Background()
+	registry := NewConnectionRegistry()
+
+	// Add some mock connections to the registry
+	conn1 := mocks.NewMockConnection(t)
+	conn2 := mocks.NewMockConnection(t)
+
+	conn1.EXPECT().ID().Return("conn1")
+	conn2.EXPECT().ID().Return("conn2")
+
+	registry.connections[conn1.ID()] = conn1
+	registry.connections[conn2.ID()] = conn2
+
+	// Set up expectations for the Close method
+	conn1.EXPECT().Close(ctx, websocket.StatusServiceRestart, "").Return(nil)
+	conn2.EXPECT().Close(ctx, websocket.StatusServiceRestart, "").Return(nil)
+
+	err := registry.Shutdown(ctx)
+
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	// Verify that the registry is closed
+	if !registry.isClosed {
+		t.Error("Expected registry to be closed")
+	}
+}