Fix session tracking test - wait for all sessions to be processed

v2/sessions/tracker_test.go

@@ -7,6 +7,8 @@ import (
 	"time"
 )
 
+const sessionsCount = 50000
+
 type testPublisher struct {
 	mutex            sync.Mutex
 	sessionsReceived [][]*Session
@@ -45,18 +47,27 @@ func TestStartSessionModifiesContext(t *testing.T) {
 func TestShouldOnlyWriteWhenReceivingSessions(t *testing.T) {
 	st, c := makeSessionTracker()
 	defer close(c)
+	wg := &sync.WaitGroup{}
+
 	go st.processSessions()
 	time.Sleep(10 * st.config.PublishInterval) // Would publish many times in this time period if there were sessions
 
 	if got := pub.sessionsReceived; len(got) != 0 {
 		t.Errorf("pub was invoked unexpectedly %d times with arguments: %v", len(got), got)
 	}
 
-	for i := 0; i < 50000; i++ {
-		st.StartSession(context.Background())
+	for i := 0; i < sessionsCount; i++ {
+		wg.Add(1)
+		go func(wg *sync.WaitGroup) {
+			defer wg.Done()
+			st.StartSession(context.Background())
+		}(wg)
 	}
 	time.Sleep(st.config.PublishInterval * 2)
 
+	// wait for all of the sessions to get consumed
+	wg.Wait()
+
 	var sessions []*Session
 	pub.mutex.Lock()
 	defer pub.mutex.Unlock()
@@ -66,7 +77,7 @@ func TestShouldOnlyWriteWhenReceivingSessions(t *testing.T) {
 			sessions = append(sessions, session)
 		}
 	}
-	if exp, got := 50000, len(sessions); exp != got {
+	if exp, got := sessionsCount, len(sessions); exp != got {
 		t.Errorf("Expected %d sessions but got %d", exp, got)
 	}