fix memory pubsub

Author: anbraten

services/pubsub/memory.go

@@ -49,14 +49,17 @@ func (p *Memory) Subscribe(c context.Context, topic string, subscriber Subscribe
 	p.topics[topic][&subscriber] = struct{}{}
 	p.Unlock()
 
-	// Wait for context to be done
-	<-c.Done()
+	// Unsubscribe when context is done
+	go func() {
+		// Wait for context to be done
+		<-c.Done()
 
-	// Unsubscribe
-	p.Lock()
-	delete(p.topics[topic], &subscriber)
-	if len(p.topics[topic]) == 0 {
-		delete(p.topics, topic)
-	}
-	p.Unlock()
+		// Unsubscribe
+		p.Lock()
+		delete(p.topics[topic], &subscriber)
+		if len(p.topics[topic]) == 0 {
+			delete(p.topics, topic)
+		}
+		p.Unlock()
+	}()
 }

services/pubsub/memory_test.go

@@ -7,7 +7,6 @@ import (
 	"context"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -24,20 +23,11 @@ func TestPubsub(t *testing.T) {
 	)
 
 	broker := NewMemory()
-	go func() {
-		broker.Subscribe(ctx, testTopic, func(message []byte) { assert.Equal(t, testMessage, message); wg.Done() })
-	}()
-	go func() {
-		broker.Subscribe(ctx, testTopic, func(_ []byte) { wg.Done() })
-	}()
-
-	// Wait a bit for the subscriptions to be registered
-	<-time.After(100 * time.Millisecond)
+	broker.Subscribe(ctx, testTopic, func(message []byte) { assert.Equal(t, testMessage, message); wg.Done() })
+	broker.Subscribe(ctx, testTopic, func(_ []byte) { wg.Done() })
 
 	wg.Add(2)
-	go func() {
-		broker.Publish(ctx, testTopic, testMessage)
-	}()
+	broker.Publish(ctx, testTopic, testMessage)
 
 	wg.Wait()
 	cancel(nil)