edge/pkg/bus/testing/request_reply.go

package testing

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

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

const (
	testTypeReqResAddress bus.Address = "testTypeReqResAddress"
)

func TestRequestReply(t *testing.T, b bus.Bus) {
	expectedRoundTrips := 256
	timeout := time.Now().Add(time.Duration(expectedRoundTrips) * time.Second)

	replyCtx, cancelReply := context.WithDeadline(context.Background(), timeout)
	defer cancelReply()

	var resWaitGroup sync.WaitGroup

	replyErrs := b.Reply(replyCtx, testTypeReqResAddress, func(env bus.Envelope) (any, error) {
		defer resWaitGroup.Done()

		req, ok := env.Message().(int)
		if !ok {
			return nil, errors.WithStack(bus.ErrUnexpectedMessage)
		}

		// Simulate random work
		time.Sleep(time.Millisecond * 100)

		t.Logf("[RES] sending res #%d", req)

		return req, nil
	})

	go func() {
		for err := range replyErrs {
			if !errors.Is(err, context.Canceled) {
				t.Errorf("%+v", errors.WithStack(err))
			}
		}
	}()

	var reqWaitGroup sync.WaitGroup

	for i := 0; i < expectedRoundTrips; i++ {
		resWaitGroup.Add(1)
		reqWaitGroup.Add(1)

		go func(i int) {
			defer reqWaitGroup.Done()

			requestCtx, cancelRequest := context.WithDeadline(context.Background(), timeout)
			defer cancelRequest()

			t.Logf("[REQ] sending req #%d", i)

			response, err := b.Request(requestCtx, bus.NewEnvelope(testTypeReqResAddress, i))
			if err != nil {
				t.Error(err)
			}

			t.Logf("[REQ] received req #%d reply", i)

			if response == nil {
				t.Error("response should not be nil")

				return
			}

			result, ok := response.Message().(int)
			if !ok {
				t.Error(errors.WithStack(bus.ErrUnexpectedMessage))

				return
			}

			if e, g := i, result; e != g {
				t.Errorf("response.Message(): expected '%v', got '%v'", e, g)
			}
		}(i)
	}

	reqWaitGroup.Wait()
	resWaitGroup.Wait()
}

func TestCanceledRequest(t *testing.T, b bus.Bus) {
	replyCtx, cancelReply := context.WithCancel(context.Background())
	defer cancelReply()

	errs := b.Reply(replyCtx, testTypeReqResAddress, func(env bus.Envelope) (any, error) {
		return env.Message(), nil
	})

	go func() {
		for err := range errs {
			if !errors.Is(err, context.Canceled) {
				t.Errorf("%+v", errors.WithStack(err))
			}
		}
	}()

	var wg sync.WaitGroup

	count := 100

	wg.Add(count)

	for i := 0; i < count; i++ {
		go func(i int) {
			defer wg.Done()

			t.Logf("calling  %d", i)

			isCanceled := i%2 == 0

			var ctx context.Context
			if isCanceled {
				canceledCtx, cancel := context.WithCancel(context.Background())
				cancel()
				ctx = canceledCtx
			} else {
				ctx = context.Background()
			}

			t.Logf("publishing envelope #%d", i)

			reply, err := b.Request(ctx, bus.NewEnvelope(testTypeReqResAddress, int64(i)))
			if err != nil {
				if errors.Is(err, context.Canceled) && isCanceled {
					return
				}

				if errors.Is(err, bus.ErrNoResponse) && isCanceled {
					return
				}

				t.Errorf("%+v", errors.WithStack(err))

				return
			}

			result, ok := reply.Message().(int64)
			if !ok {
				t.Errorf("response.Result: expected type '%T', got '%T'", int64(0), reply.Message())

				return
			}

			if e, g := i, int(result); e != g {
				t.Errorf("response.Result: expected '%v', got '%v'", e, g)

				return
			}
		}(i)
	}

	wg.Wait()
}