Merge pull request #3 from andrewwormald/andreww-allowAccessToPoolKeys

allow access to connection pool keys

mock/mock.go

@@ -22,7 +22,7 @@ type Acceptor interface {
 	Accept(w http.ResponseWriter, r *http.Request, channelKey string) error
 }
 
-// NewChannelPool returns a pool of channel side ws connections that get injected via the Accept method which Stream
+// NewChannelPool returns a pool of client side ws connections that get injected via the Accept method which Stream
 // satisfies.
 func NewChannelPool(t *testing.T, a Acceptor, connectionCount int) []WSChannel {
 	var pool []WSChannel

stream.go

@@ -109,13 +109,35 @@ func (s *Stream) Read() chan ReceiveMessage {
 }
 
 // Connections returns the amount of valid channels that are in the Stream.
+// Deprecated: Will be removed in the future. Instead, use ConnectionsCount.
 func (s *Stream) Connections() int {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	return len(s.pool)
 }
 
+// ConnectionsCount returns the amount of valid channels that are in the Stream.
+func (s *Stream) ConnectionsCount() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	return len(s.pool)
+}
+
+// ConnectionKeys returns the keys of all the current valid connections.
+func (s *Stream) ConnectionKeys() []string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var keys []string
+	for key := range s.pool {
+		keys = append(keys, key)
+	}
+
+	return keys
+}
+
 // store safely adds the channel to the Stream's channel pool without causing any data races.
 func (s *Stream) store(c *Channel) {
 	s.mu.Lock()

stream_test.go

@@ -99,4 +99,14 @@ func TestConnectionsCount(t *testing.T) {
 	connectionCount := 1000
 	_ = mock.NewChannelPool(t, s, connectionCount)
 	require.Equal(t, connectionCount, s.Connections())
+	require.Equal(t, connectionCount, s.ConnectionsCount())
+}
+
+func TestConnectionsKeys(t *testing.T) {
+	ctx := context.TODO()
+	s := New(ctx)
+	connectionCount := 1000
+	_ = mock.NewChannelPool(t, s, connectionCount)
+
+	require.Equal(t, connectionCount, len(s.ConnectionKeys()))
 }