edge/pkg/bus/testing/publish_subscribe.go

package testing

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"forge.cadoles.com/arcad/edge/pkg/bus"
	"github.com/pkg/errors"
)

const (
	testAddress bus.Address = "testAddress"
)

func TestPublishSubscribe(t *testing.T, b bus.Bus) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
	defer cancel()

	t.Log("subscribe")

	envelopes, err := b.Subscribe(ctx, testAddress)
	if err != nil {
		t.Fatal(errors.WithStack(err))
	}

	expectedTotal := 5

	var wg sync.WaitGroup

	wg.Add(expectedTotal)

	go func() {

		count := expectedTotal

		for i := 0; i < count; i++ {
			env := bus.NewEnvelope(testAddress, fmt.Sprintf("message %d", i))

			if err := b.Publish(env); err != nil {
				t.Error(errors.WithStack(err))
			}

			t.Logf("published %d", i)
		}
	}()

	var count int32 = 0

	go func() {
		t.Log("range for received envelopes")

		for env := range envelopes {
			t.Logf("received msg %d", atomic.LoadInt32(&count))
			atomic.AddInt32(&count, 1)

			if e, g := testAddress, env.Address(); e != g {
				t.Errorf("env.Address(): expected '%v', got '%v'", e, g)
			}

			wg.Done()
		}
	}()

	wg.Wait()

	b.Unsubscribe(testAddress, envelopes)

	if e, g := int32(expectedTotal), count; e != g {
		t.Errorf("envelopes received count: expected '%v', got '%v'", e, g)
	}
}