From d11d48d23d62c890b382b053e199ac81eea621a7 Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <k.sysoev@me.com>
Date: Sat, 27 Apr 2024 15:13:34 +0800
Subject: [PATCH] Adds tests for closing channel and connection registry

---
 channel/channel_test.go             | 14 ++++++++++++++
 channel/connection_registry.go      |  2 +-
 channel/connection_registry_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/channel/channel_test.go b/channel/channel_test.go
index 52c6d8a..9b96174 100644
--- a/channel/channel_test.go
+++ b/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)
+	}
+}
diff --git a/channel/connection_registry.go b/channel/connection_registry.go
index 602a36a..46b6bd2 100644
--- a/channel/connection_registry.go
+++ b/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()
 		}()
 	}
 
diff --git a/channel/connection_registry_test.go b/channel/connection_registry_test.go
index 10f86d4..f3da683 100644
--- a/channel/connection_registry_test.go
+++ b/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")
+	}
+}