summaryrefslogtreecommitdiffhomepage
path: root/internal/pubsub/broker_test.go
blob: b4caa98f3d136c79bd28c89a6210a2a03b8a84dc (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
package pubsub

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestBrokerSubscribe(t *testing.T) {
	t.Parallel()

	t.Run("with cancellable context", func(t *testing.T) {
		t.Parallel()
		broker := NewBroker[string]()
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		ch := broker.Subscribe(ctx)
		assert.NotNil(t, ch)
		assert.Equal(t, 1, broker.GetSubscriberCount())

		// Cancel the context should remove the subscription
		cancel()
		time.Sleep(10 * time.Millisecond) // Give time for goroutine to process
		assert.Equal(t, 0, broker.GetSubscriberCount())
	})

	t.Run("with background context", func(t *testing.T) {
		t.Parallel()
		broker := NewBroker[string]()

		// Using context.Background() should not leak goroutines
		ch := broker.Subscribe(context.Background())
		assert.NotNil(t, ch)
		assert.Equal(t, 1, broker.GetSubscriberCount())

		// Shutdown should clean up all subscriptions
		broker.Shutdown()
		assert.Equal(t, 0, broker.GetSubscriberCount())
	})
}

func TestBrokerPublish(t *testing.T) {
	t.Parallel()
	broker := NewBroker[string]()
	ctx := t.Context()

	ch := broker.Subscribe(ctx)

	// Publish a message
	broker.Publish(EventTypeCreated, "test message")

	// Verify message is received
	select {
	case event := <-ch:
		assert.Equal(t, EventTypeCreated, event.Type)
		assert.Equal(t, "test message", event.Payload)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout waiting for message")
	}
}

func TestBrokerShutdown(t *testing.T) {
	t.Parallel()
	broker := NewBroker[string]()

	// Create multiple subscribers
	ch1 := broker.Subscribe(context.Background())
	ch2 := broker.Subscribe(context.Background())

	assert.Equal(t, 2, broker.GetSubscriberCount())

	// Shutdown should close all channels and clean up
	broker.Shutdown()

	// Verify channels are closed
	_, ok1 := <-ch1
	_, ok2 := <-ch2
	assert.False(t, ok1, "channel 1 should be closed")
	assert.False(t, ok2, "channel 2 should be closed")

	// Verify subscriber count is reset
	assert.Equal(t, 0, broker.GetSubscriberCount())
}

func TestBrokerConcurrency(t *testing.T) {
	t.Parallel()
	broker := NewBroker[int]()

	// Create a large number of subscribers
	const numSubscribers = 100
	var wg sync.WaitGroup
	wg.Add(numSubscribers)

	// Create a channel to collect received events
	receivedEvents := make(chan int, numSubscribers)

	for i := range numSubscribers {
		go func(id int) {
			defer wg.Done()
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()

			ch := broker.Subscribe(ctx)

			// Receive one message then cancel
			select {
			case event := <-ch:
				receivedEvents <- event.Payload
			case <-time.After(1 * time.Second):
				t.Errorf("timeout waiting for message %d", id)
			}
			cancel()
		}(i)
	}

	// Give subscribers time to set up
	time.Sleep(10 * time.Millisecond)

	// Publish messages to all subscribers
	for i := range numSubscribers {
		broker.Publish(EventTypeCreated, i)
	}

	// Wait for all subscribers to finish
	wg.Wait()
	close(receivedEvents)

	// Give time for cleanup goroutines to run
	time.Sleep(10 * time.Millisecond)

	// Verify all subscribers are cleaned up
	assert.Equal(t, 0, broker.GetSubscriberCount())

	// Verify we received the expected number of events
	count := 0
	for range receivedEvents {
		count++
	}
	assert.Equal(t, numSubscribers, count)
}