refactor: cleanup code

Author: Tochemey

eventstream/stream.go

@@ -94,16 +94,16 @@ func (b *EventsStream) RemoveSubscriber(sub Subscriber) {
 
 // Broadcast notifies all subscribers of a given topic of a new message
 func (b *EventsStream) Broadcast(msg any, topics []string) {
-	// broadcast message to all topics.
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
 	for _, topic := range topics {
 		for _, consumer := range b.topics[topic] {
-			m := NewMessage(topic, msg)
 			if !consumer.Active() {
-				return
+				continue
 			}
-			go (func(s Subscriber) {
-				s.signal(m)
-			})(consumer)
+			m := NewMessage(topic, msg)
+			go consumer.signal(m)
 		}
 	}
 }
@@ -136,33 +136,31 @@ func (b *EventsStream) Subscribe(sub Subscriber, topic string) {
 
 // Unsubscribe removes a subscriber from a topic
 func (b *EventsStream) Unsubscribe(sub Subscriber, topic string) {
-	// unsubscribe to given topic
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
-	// only unsubscribe active subscriber
-	if !sub.Active() {
-		return
+	if subscribers, ok := b.topics[topic]; ok {
+		delete(subscribers, sub.ID())
+		if len(subscribers) == 0 {
+			delete(b.topics, topic)
+		}
 	}
 
-	delete(b.topics[topic], sub.ID())
 	sub.unsubscribe(topic)
 }
 
 // Publish publishes a message to a topic
 func (b *EventsStream) Publish(topic string, msg any) {
-	// publish the message to given topic.
 	b.mu.Lock()
-	bTopics := b.topics[topic]
+	subscribers := b.topics[topic]
 	b.mu.Unlock()
-	for _, consumer := range bTopics {
-		m := NewMessage(topic, msg)
+
+	for _, consumer := range subscribers {
 		if !consumer.Active() {
-			return
+			continue
 		}
-		go (func(s Subscriber) {
-			s.signal(m)
-		})(consumer)
+		m := NewMessage(topic, msg)
+		go consumer.signal(m)
 	}
 }
 

eventstream/stream_test.go

@@ -53,9 +53,7 @@ func TestStream(t *testing.T) {
 		broker.Subscribe(cons, "t3")
 		assert.Zero(t, broker.SubscribersCount("t3"))
 
-		t.Cleanup(func() {
-			broker.Close()
-		})
+		broker.Close()
 	})
 	t.Run("With Unsubscription", func(t *testing.T) {
 		broker := New()
@@ -66,12 +64,19 @@ func TestStream(t *testing.T) {
 		broker.Subscribe(cons, "t1")
 		broker.Subscribe(cons, "t2")
 
-		require.EqualValues(t, 1, broker.SubscribersCount("t1"))
+		sub2 := broker.AddSubscriber()
+		require.NotNil(t, sub2)
+		broker.Subscribe(sub2, "t1")
+
+		require.EqualValues(t, 2, broker.SubscribersCount("t1"))
 		require.EqualValues(t, 1, broker.SubscribersCount("t2"))
 
+		sub2.Shutdown()
+
 		// Unsubscribe the consumer
 		broker.Unsubscribe(cons, "t1")
-		assert.Zero(t, broker.SubscribersCount("t1"))
+		broker.Unsubscribe(sub2, "t1")
+		require.Zero(t, broker.SubscribersCount("t1"))
 		require.EqualValues(t, 1, broker.SubscribersCount("t2"))
 
 		broker.Subscribe(cons, "t3")
@@ -82,9 +87,7 @@ func TestStream(t *testing.T) {
 		broker.Subscribe(cons, "t4")
 		assert.Zero(t, broker.SubscribersCount("t4"))
 
-		t.Cleanup(func() {
-			broker.Close()
-		})
+		broker.Close()
 	})
 	t.Run("With Publication", func(t *testing.T) {
 		broker := New()
@@ -95,25 +98,32 @@ func TestStream(t *testing.T) {
 		broker.Subscribe(cons, "t1")
 		broker.Subscribe(cons, "t2")
 
-		require.EqualValues(t, 1, broker.SubscribersCount("t1"))
+		sub2 := broker.AddSubscriber()
+		require.NotNil(t, sub2)
+		broker.Subscribe(sub2, "t1")
+
+		require.EqualValues(t, 2, broker.SubscribersCount("t1"))
 		require.EqualValues(t, 1, broker.SubscribersCount("t2"))
 
+		sub2.Shutdown()
+
 		broker.Publish("t1", "hi")
 		broker.Publish("t2", "hello")
 
 		time.Sleep(time.Second)
 
 		var messages []*Message
 		for message := range cons.Iterator() {
+			require.NotNil(t, message)
+			require.NotNil(t, message.Topic())
+			require.NotNil(t, message.Payload())
 			messages = append(messages, message)
 		}
 
 		assert.Len(t, messages, 2)
 		assert.Len(t, cons.Topics(), 2)
 
-		t.Cleanup(func() {
-			broker.Close()
-		})
+		broker.Close()
 	})
 	t.Run("With Broadcast", func(t *testing.T) {
 		broker := New()
@@ -124,9 +134,15 @@ func TestStream(t *testing.T) {
 		broker.Subscribe(cons, "t1")
 		broker.Subscribe(cons, "t2")
 
-		require.EqualValues(t, 1, broker.SubscribersCount("t1"))
+		sub2 := broker.AddSubscriber()
+		require.NotNil(t, sub2)
+		broker.Subscribe(sub2, "t1")
+
+		require.EqualValues(t, 2, broker.SubscribersCount("t1"))
 		require.EqualValues(t, 1, broker.SubscribersCount("t2"))
 
+		sub2.Shutdown()
+
 		broker.Broadcast("hi", []string{"t1", "t2"})
 
 		time.Sleep(time.Second)
@@ -139,8 +155,6 @@ func TestStream(t *testing.T) {
 		assert.Len(t, messages, 2)
 		assert.Len(t, cons.Topics(), 2)
 
-		t.Cleanup(func() {
-			broker.Close()
-		})
+		broker.Close()
 	})
 }