path: root/fanout/fanout_test.go

package fanout

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

type mockObserver struct {
	name string
	t    *testing.T
}

func (m *mockObserver) OnDistribute() {
	m.t.Logf("fanout: %s OnDistribute", m.name)
}

func (m *mockObserver) OnOutputClosed() {
	m.t.Logf("fanout: %s OnOutputClosed", m.name)
}

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

	tests := []struct {
		name      string
		opts      []Option
		shouldErr bool
	}{
		{
			name:      "worker size is invalid",
			opts:      []Option{WithWorkerCount(-10)},
			shouldErr: true,
		},
		{
			name:      "buffer size is invalid",
			opts:      []Option{WithBufferSize(-10), WithWorkerCount(1)},
			shouldErr: true,
		},
		{
			name:      "valid options",
			opts:      []Option{WithBufferSize(1), WithWorkerCount(1)},
			shouldErr: false,
		},
		{
			name:      "worker count is zero",
			opts:      []Option{WithWorkerCount(0)},
			shouldErr: true,
		},
		{
			name:      "valid broadcast strategy",
			opts:      []Option{WithWorkerCount(5), WithBufferSize(10), WithStrategy(Broadcast)},
			shouldErr: false,
		},
		{
			name:      "invalid strategy",
			opts:      []Option{WithWorkerCount(1), WithStrategy(FanOutStrategy(99))},
			shouldErr: true,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()
			_, err := NewFanOut[int](test.opts...)
			if (err != nil) != test.shouldErr {
				t.Errorf("validation error = %v, want err %v", err, test.shouldErr)
			}
		})
	}
}

func TestFanOut_ConcurrentStress(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()

	workerCount := 5
	fo, err := NewFanOut[int](
		WithWorkerCount(workerCount),
		WithBufferSize(100),
	)
	if err != nil {
		t.Fatalf("new fanout error: %v", err)
	}

	input := make(chan int, 1000)
	fo.Run(ctx, input)

	var wg sync.WaitGroup

	for i := 0; i < 10; i++ {
		wg.Go(func() {
			outs := fo.Outputs()
			if len(outs) != workerCount {
				t.Errorf("expected %d outputs, got %d", workerCount, len(outs))
			}
		})
	}

	for _, ch := range fo.Outputs() {
		wg.Go(func() {
			for range ch {
			}
		})
	}

	producers := 10
	messagesPerProducer := 100

	var producerWg sync.WaitGroup
	producerWg.Add(producers)

	for i := 0; i < producers; i++ {
		go func() {
			defer producerWg.Done()
			for j := 0; j < messagesPerProducer; j++ {
				select {
				case input <- j:
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	go func() {
		producerWg.Wait()
		close(input)
	}()

	wg.Wait()
}

func TestFanOut_RunBroadcast(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup
	workerCount := 2
	size := 10
	obs := &mockObserver{
		name: "fo broadcast",
		t:    t,
	}
	fo, _ := NewFanOut[int](
		WithWorkerCount(workerCount),
		WithStrategy(Broadcast),
		WithObserver(obs),
	)

	ch := make(chan int, size)
	wg.Go(func() {
		for i := range size {
			ch <- i
		}
		close(ch)
	})

	fo.Run(t.Context(), ch)

	res := make([]int, workerCount)
	for i, cha := range fo.Outputs() {
		wg.Go(func() {
			count := 0
			for v := range cha {
				t.Logf("%v", v)
				count++
				res[i] = count
			}
		})
	}

	wg.Wait()

	expected := size
	for i, count := range res {
		if count != expected {
			t.Errorf("worker %d received %d items, expected %d", i, count, expected)
		}
	}
}

func TestFanOut_RunRoundRobin(t *testing.T) {
	t.Parallel()
	var wg sync.WaitGroup
	size := 10
	workerCount := 2
	obs := &mockObserver{
		name: "fo round robin",
		t:    t,
	}
	fo, _ := NewFanOut[int](
		WithWorkerCount(workerCount),
		WithStrategy(RoundRobin),
		WithObserver(obs),
	)

	ch := make(chan int, size)
	wg.Go(func() {
		for i := range size {
			ch <- i
		}
		close(ch)
	})

	fo.Run(t.Context(), ch)
	res := make([]int, workerCount)
	for i, cha := range fo.Outputs() {
		wg.Go(func() {
			count := 0
			for v := range cha {
				t.Logf("%v", v)
				count++
				res[i] = count
			}
		})
	}

	wg.Wait()
	expected := size / workerCount

	for i, count := range res {
		if count != expected {
			t.Errorf("worker %d received %d items, expected %d", i, count, expected)
		}
	}
}

func TestFanOut_ContextTimeOut(t *testing.T) {
	t.Parallel()
	fo, _ := NewFanOut[int](WithWorkerCount(2))
	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
	defer cancel()

	input := make(chan int, 10)
	fo.Run(ctx, input)

	var wg sync.WaitGroup
	done := make(chan struct{})
	for _, ch := range fo.Outputs() {
		wg.Go(func() {
			for v := range ch {
				t.Logf("received %d", v)
			}
		})
	}

	go func() {
		wg.Wait()
		close(done)
	}()

	go func() {
		i := 0
		for {
			select {
			case input <- i:
				i++
				time.Sleep(time.Millisecond)
			case <-ctx.Done():
				return
			}
		}
	}()

	select {
	case <-done:
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout waiting for cleanup")
	}
}

func TestFanOut_ContextCancel(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(t.Context())

	fo, err := NewFanOut[int](WithWorkerCount(3))
	if err != nil {
		t.Fatalf("failed to create fanout: %v", err)
	}

	input := make(chan int, 10)
	fo.Run(ctx, input)

	var wg sync.WaitGroup
	for _, ch := range fo.Outputs() {
		wg.Go(func() {
			for range ch {
			}
		})
	}

	cancel()
	close(input)

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("timeout waiting for shutdown")
	}
}

func TestFanOut_ContextDoneRoundRobin(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(t.Context())

	fo, err := NewFanOut[int](
		WithWorkerCount(3),
		WithBufferSize(0),
		WithStrategy(RoundRobin),
	)
	if err != nil {
		t.Fatalf("new fanout error: %v", err)
	}

	input := make(chan int, 1)
	input <- 1
	fo.Run(ctx, input)

	cancel()

	for i, ch := range fo.Outputs() {
		select {
		case _, ok := <-ch:
			if ok {
				for range ch {
				}
			}
		case <-time.After(time.Second):
			t.Fatalf("worker %d: timeout waiting for close", i)
		}
	}
}

func TestFanOut_ContextDoneBroadcast(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(t.Context())

	fo, err := NewFanOut[int](
		WithWorkerCount(3),
		WithStrategy(Broadcast),
		WithBufferSize(0),
	)
	if err != nil {
		t.Fatalf("new fanout error: %v", err)
	}

	input := make(chan int, 1)
	input <- 1
	fo.Run(ctx, input)

	cancel()

	for i, ch := range fo.Outputs() {
		select {
		case _, ok := <-ch:
			if ok {
				for range ch {
				}
			}
		case <-time.After(time.Second):
			t.Fatalf("worker %d: timeout waiting for close", i)
		}
	}
}

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

	ctx := t.Context()
	fo, err := NewFanOut[int](
		WithWorkerCount(3),
		WithBufferSize(0),
	)
	if err != nil {
		t.Fatalf("new fanout error: %v", err)
	}

	input := make(chan int, 10)
	fo.Run(ctx, input)

	defer func() {
		if r := recover(); r == nil {
			t.Error("expected panic when calling Run() multiple times")
		}
	}()

	fo.Run(ctx, input)
}

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

	fo, err := NewFanOut[int](WithWorkerCount(1))
	if err != nil {
		t.Fatalf("failed to create fanout: %v", err)
	}

	defer func() {
		if r := recover(); r == nil {
			t.Error("expected panic when passing nil input channel")
		}
	}()

	fo.Run(t.Context(), nil)
}