completed stream request flow and tests. restructured other parts of the code.

postmaster_test.go

@@ -2,79 +2,101 @@ package prism
 
 import (
 	"context"
+	"fmt"
 	"github.com/stretchr/testify/assert"
 	"testing"
 	"time"
 )
 
-const testURL = "wss://test"
+// Helpers
 
-func TestPostmasterUnknownPeerSend(t *testing.T) {
-	ctx := context.Background()
+func mockPostmaster(
+	ctx context.Context,
+) (pm *Postmaster, poolEvents chan PoolEvent) {
 	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
+	poolEvents = make(chan PoolEvent, 4)
 	poolSendFunc := func(id string, data Envelope) error { return nil }
+	pm = NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
+	return
+}
 
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
-
-	called := make(chan LetterOutcome, 1)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+func expectLetterOutcome(
+	t *testing.T, ch chan LetterOutcome, kind LetterOutcomeKind,
+) {
+	t.Helper()
 
 	var outcome LetterOutcome
 	Eventually(t, func() bool {
 		select {
 		default:
 			return false
-		case outcome = <-called:
+		case outcome = <-ch:
 			return true
 		}
 	}, "should have received outcome")
 
-	assert.Equal(t, OutcomeRejected, outcome.Kind)
+	assert.Equal(t, kind, outcome.Kind)
+}
+
+func expectAllLetterOutcomes(
+	t *testing.T, ch chan LetterOutcome, kind LetterOutcomeKind, count int,
+) {
+	t.Helper()
+
+	outcomes := make([]LetterOutcome, 0, count)
+	Eventually(t, func() bool {
+		select {
+		default:
+			return false
+		case o := <-ch:
+			outcomes = append(outcomes, o)
+			return len(outcomes) == count
+		}
+	}, fmt.Sprintf("should have returned %d outcomes", count))
+
+	if len(outcomes) >= count {
+		for i := range count {
+			assert.Equal(t, OutcomeCancelled, outcomes[i].Kind)
+		}
+	}
+}
+
+// Tests
+
+func TestPostmasterUnknownPeerSend(t *testing.T) {
+	ctx := context.Background()
+	pm, _ := mockPostmaster(ctx)
+
+	called := make(chan LetterOutcome, 1)
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
+
+	expectLetterOutcome(t, called, OutcomeRejected)
 }
 
 func TestPostmasterSend(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
+	pm, poolEvents := mockPostmaster(ctx)
 
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
-
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventConnected, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventConnected, At: time.Now()}
 
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	called := make(chan LetterOutcome, 1)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
 
-	var outcome LetterOutcome
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case outcome = <-called:
-			return true
-		}
-	}, "should have received outcome")
-
-	assert.Equal(t, OutcomeSent, outcome.Kind)
+	expectLetterOutcome(t, called, OutcomeSent)
 }
 
 func TestPostmasterCancelInFlight(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
+	pm, poolEvents := mockPostmaster(ctx)
 
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
-
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	called := make(chan LetterOutcome, 1)
-	cancel := pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+	_, cancel := pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
 
 	// wait for letter to queue
 	time.Sleep(100 * time.Millisecond)
@@ -83,133 +105,74 @@ func TestPostmasterCancelInFlight(t *testing.T) {
 	cancel()
 
 	// connect the pool
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventConnected, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventConnected, At: time.Now()}
 
-	var outcome LetterOutcome
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case outcome = <-called:
-			return true
-		}
-	}, "should have received outcome")
-
-	// letter should drain out of the queue and return cancelled
-	assert.Equal(t, OutcomeCancelled, outcome.Kind)
+	expectLetterOutcome(t, called, OutcomeCancelled)
 }
 
 func TestPostmasterExpire(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
+	pm, poolEvents := mockPostmaster(ctx)
 
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
-
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	called := make(chan LetterOutcome, 1)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o },
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o },
 		WithDeadline(1*time.Millisecond))
 
 	// wait for letter to queue and expire
 	time.Sleep(100 * time.Millisecond)
 
 	// connect the pool
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventConnected, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventConnected, At: time.Now()}
 
-	var outcome LetterOutcome
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case outcome = <-called:
-			return true
-		}
-	}, "should have received outcome")
-
-	// letter should drain out of the queue and return expired
-	assert.Equal(t, OutcomeExpired, outcome.Kind)
+	expectLetterOutcome(t, called, OutcomeExpired)
 }
 
 func TestPostmasterPeerRemoved(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
-
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
+	pm, poolEvents := mockPostmaster(ctx)
 
 	// add peer, but do not connect
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	// send two letters
-	outcomes := make([]LetterOutcome, 0, 2)
 	called := make(chan LetterOutcome, 2)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
 
 	// wait for them to hit the courier queue
 	time.Sleep(100 * time.Millisecond)
 
 	// remove the peer
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventRemoved, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventRemoved, At: time.Now()}
 
 	// expect each letter to return cancelled
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case o := <-called:
-			outcomes = append(outcomes, o)
-			return len(outcomes) == 2
-		}
-	}, "should have returned 2 outcomes")
-
-	if len(outcomes) >= 2 {
-		assert.Equal(t, OutcomeCancelled, outcomes[0].Kind)
-		assert.Equal(t, OutcomeCancelled, outcomes[1].Kind)
-	}
+	expectAllLetterOutcomes(t, called, OutcomeCancelled, 2)
 
 	// subsequent sends should fail
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
-
-	var outcome LetterOutcome
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case outcome = <-called:
-			return true
-		}
-	}, "should have received outcome")
-
-	assert.Equal(t, OutcomeRejected, outcome.Kind)
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
+	expectLetterOutcome(t, called, OutcomeRejected)
 }
 
 func TestPostmasterCourierCloseRace(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
-
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
+	pm, poolEvents := mockPostmaster(ctx)
 
 	// add peer, but do not connect
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	// remove the peer
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventRemoved, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventRemoved, At: time.Now()}
 
 	// send a letter
 	time.Sleep(5 * time.Microsecond) // small wait lines up the race condition
 	var outcome *LetterOutcome
 	called := make(chan LetterOutcome, 1)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
 
 	Eventually(t, func() bool {
 		select {
@@ -236,21 +199,16 @@ func TestPostmasterCourierCloseRace(t *testing.T) {
 
 func TestPostmasterClose(t *testing.T) {
 	ctx := context.Background()
-	poolHasPeer := func(id string) (string, bool) { return id, true }
-	poolEvents := make(chan PoolEvent, 4)
-	poolSendFunc := func(id string, data Envelope) error { return nil }
-
-	pm := NewPostmaster(ctx, poolHasPeer, poolEvents, poolSendFunc, nil)
+	pm, poolEvents := mockPostmaster(ctx)
 
 	// add peer, but do not connect
-	poolEvents <- PoolEvent{ID: testURL, Kind: EventAdded, At: time.Now()}
+	poolEvents <- PoolEvent{ID: "peer", Kind: EventAdded, At: time.Now()}
 	Eventually(t, func() bool { return len(pm.Peers()) > 0 }, "should add peer")
 
 	// send two letters
-	outcomes := make([]LetterOutcome, 0, 2)
 	called := make(chan LetterOutcome, 2)
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
 
 	// wait for them to hit the courier queue
 	time.Sleep(100 * time.Millisecond)
@@ -259,33 +217,9 @@ func TestPostmasterClose(t *testing.T) {
 	pm.Close()
 
 	// expect each letter to return cancelled
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case o := <-called:
-			outcomes = append(outcomes, o)
-			return len(outcomes) == 2
-		}
-	}, "should have returned 2 outcomes")
-
-	if len(outcomes) >= 2 {
-		assert.Equal(t, OutcomeCancelled, outcomes[0].Kind)
-		assert.Equal(t, OutcomeCancelled, outcomes[1].Kind)
-	}
+	expectAllLetterOutcomes(t, called, OutcomeCancelled, 2)
 
 	// subsequent sends should be rejected
-	pm.Send(ctx, testURL, nil, func(o LetterOutcome) { called <- o })
-
-	var outcome LetterOutcome
-	Eventually(t, func() bool {
-		select {
-		default:
-			return false
-		case outcome = <-called:
-			return true
-		}
-	}, "should have received outcome")
-
-	assert.Equal(t, OutcomeRejected, outcome.Kind)
+	pm.Send(ctx, "peer", nil, func(o LetterOutcome) { called <- o })
+	expectLetterOutcome(t, called, OutcomeRejected)
 }