taskgroup_test.go

package taskgroup_test

import (
	"context"
	"fmt"
	"testing"
	"time"

	tg "github.com/weirdname404/taskgroup"
)

const (
	N            = 10
	TestTaskTime = 100
)

var ERROR = fmt.Errorf("ERROR")
var ExpectedRes []int

type HasFatal interface {
	Fatal(args ...any)
	Fatalf(format string, args ...any)
}

func testFunc(i int) (int, error) {
	time.Sleep((time.Duration(TestTaskTime-i*10) * time.Millisecond))
	return i, nil
}

func errFunc(i int) (int, error) {
	time.Sleep((time.Duration(TestTaskTime-i*10) * time.Millisecond))
	return 0, ERROR
}

func addNTestFunc(g *tg.TaskGroup[int]) error {
	for i := range N {
		f := func() (int, error) {
			return testFunc(i)
		}
		err := g.Add(f)
		if err != nil {
			return err
		}
	}
	return nil
}

func addTasksAndGetResults(g *tg.TaskGroup[int]) ([]int, error) {
	err := addNTestFunc(g)
	if err != nil {
		return nil, err
	}
	res := make([]int, 0, len(ExpectedRes))
	for v := range g.Start() {
		res = append(res, v)
	}
	return res, nil
}

func AddTasksAndCheckResults(g *tg.TaskGroup[int], t HasFatal) {
	res, err := addTasksAndGetResults(g)
	if err != nil {
		t.Fatalf("Unexpected error - %v", err)
	}
	if err := g.Error(); err != nil {
		t.Fatalf("Unexpected error - %v", err)
	}
	if !(len(res) == len(ExpectedRes)) {
		t.Fatalf("Unexpected length of the results - expected: %d got: %d", len(ExpectedRes), len(res))
	}
	for i := range res {
		if !(ExpectedRes[i] == res[i]) {
			t.Fatalf("Unexpected result, expected: %v got: %v", ExpectedRes[i], res[i])
		}
	}
}

func ErrorScenario(g *tg.TaskGroup[int]) ([]int, error) {
	for i := range N {
		f := func(ctx context.Context) (int, error) {
			return testFunc(i)
		}
		if i%2 == 0 {
			err := g.Add(func(ctx context.Context) (int, error) {
				return errFunc(i)
			})
			if err != nil {
				return nil, err
			}
		} else {
			err := g.Add(f)
			if err != nil {
				return nil, err
			}
		}
	}
	res := make([]int, 0, len(ExpectedRes))
	for v := range g.Start() {
		res = append(res, v)
	}
	return res, nil
}

func TestTaskGroupNoTasks(t *testing.T) {
	g := tg.NewTaskGroup[any](context.Background())
	res := make([]any, 0)
	ch := g.Start()
	for v := range ch {
		res = append(res, v)
	}
	if !(len(res) == 0) {
		t.Fatalf("Unexpected results")
	}
	_, open := <-ch
	if open {
		t.Fatalf("Expected result channel to be closed")
	}
}

func TestTaskGroup(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background())
	AddTasksAndCheckResults(g, t)
}

func TestTaskGroupRemainder(t *testing.T) {
	// N = 10 - number of tasks
	// m = 3 - chunks size
	// 10 % 3 = 1 - remainder
	g := tg.NewTaskGroup[int](context.Background()).Limit(3)
	AddTasksAndCheckResults(g, t)
}

func TestNoLimitTaskGroup(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background()).Limit(0)
	AddTasksAndCheckResults(g, t)
}

func TestLimitIsGreaterThanNTasks(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background()).Limit(N * 10)
	AddTasksAndCheckResults(g, t)
}

func TestTaskGroupErr(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background())
	res, err := ErrorScenario(g)
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	gErr := g.Error()
	if !(gErr != nil) {
		t.Fatalf("Expected error but got nothing")
	}
	if gErr.Error() != ERROR.Error() {
		t.Fatalf("Unexpected error - expected: %s got: %s", ERROR.Error(), gErr.Error())
	}
	// error in the first chunk -> no chunks -> no results
	if len(res) != 0 {
		t.Fatalf("Unexpected result, expected: nothing got: %v", res)
	}
	// calling g.Error again will return nothing
	if !(g.Error() == nil) {
		t.Fatalf("Expected nothing")
	}
}

func TestTaskGroupAutoClear(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background()).Limit(5)
	AddTasksAndCheckResults(g, t)
	// it is expected that the group will auto clear itself after successful processing
	AddTasksAndCheckResults(g, t)
	// by the way `Clear` is idempotent, so you can call it anytime
	g.Clear()
	AddTasksAndCheckResults(g, t)
}

func TestTaskGroupManualClear(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background()).Limit(5)
	// it can be a case when something went wrong while adding tasks
	for i := range N {
		f := func(ctx context.Context) (int, error) {
			return testFunc(i)
		}
		// something is happening...
		if i+1 == N {
			// at this point let's pretend some error occured that makes futher addition of the task pointless
			break
		}
		err := g.Add(f)
		if !(err == nil) {
			t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
		}
	}
	// if the user forgets to clear the group and tries to add tasks again, the further results might be unexpected
	err := addNTestFunc(g)
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	// old tasks are present
	res := make([]int, 0, 20)
	for v := range g.Start() {
		res = append(res, v)
	}
	if !(len(res) > len(ExpectedRes)) {
		t.Fatalf("Unexpected amount of tasks, expected: %v got: %v", len(ExpectedRes), len(res))
	}
	// g.Start is pointless as the number of the results will be different from the expected one
	// the group must be cleared
	g.Clear()
	AddTasksAndCheckResults(g, t)
}

func TestCustomTaskGroupCtxCancel(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), TestTaskTime*(2*time.Millisecond))
	defer cancel()
	// process tasks sequentially
	g := tg.NewTaskGroup[int](ctx).Limit(1)
	res, err := addTasksAndGetResults(g)
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	// context error goes silently
	if !(g.Error() != nil) {
		t.Fatalf("Expected error - got nothing")
	}
	// 200ms is time for 2 tasks
	if len(res) != 2 {
		t.Fatalf("Unexpected result, expected: exactly 2 results got: %v", res)
	}
}

func TestErrUnknownTaskType(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background())
	// func with no target result
	err := g.Add(func() any { return nil })
	if !(err == tg.ErrUnknownTaskType) {
		t.Fatal("Expected tg.ErrUnknownTaskType")
	}
	// func with no result
	err = g.Add(func() {})
	if !(err == tg.ErrUnknownTaskType) {
		t.Fatal("Expected tg.ErrUnknownTaskType")
	}
	// func with unknown args
	err = g.Add(func(_ any) {})
	if !(err == tg.ErrUnknownTaskType) {
		t.Fatal("Expected tg.ErrUnknownTaskType")
	}
	// func with unknown results
	err = g.Add(func() (int, error, any) { return 0, nil, nil })
	if !(err == tg.ErrUnknownTaskType) {
		t.Fatal("Expected tg.ErrUnknownTaskType")
	}
	// nil
	err = g.Add(nil)
	if !(err == tg.ErrUnknownTaskType) {
		t.Fatal("Expected ErrUnknownTaskType")
	}
}

func TestTaskType(t *testing.T) {
	g := tg.NewTaskGroup[int](context.Background())
	// func with target result
	err := g.Add(func() int { return 0 })
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	// func with target result and error
	err = g.Add(func() (int, error) { return 0, nil })
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	// func with target result and context
	err = g.Add(func(_ context.Context) int { return 0 })
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
	// func with target result, error and context
	err = g.Add(func(_ context.Context) (int, error) { return 0, nil })
	if !(err == nil) {
		t.Fatalf("Unexpected error - expected: nothing got: %s", err.Error())
	}
}

// go test -bench=. -benchmem .
func Benchmark(b *testing.B) {
	for _, bechN := range []int{100, 1_000, 10_000} {
		for _, limit := range []int{bechN / 4, bechN / 3, 10_000} {
			g := tg.NewTaskGroup[int](context.Background()).Limit(limit)
			tasks := make([]func() (int, error), bechN)
			for i := range bechN {
				f := func() (int, error) {
					return testFunc(i)
				}
				tasks[i] = f
			}
			b.ResetTimer()
			b.ReportAllocs()
			b.Run(fmt.Sprintf("tasks=%d,limit=%d", bechN, limit), func(b *testing.B) {
				for i := 0; i < b.N; i++ {
					for _, task := range tasks {
						err := g.Add(task)
						if err != nil {
							b.Fatal(err)
						}
					}
					for v := range g.Start() {
						_ = v
					}
				}
			})
		}
	}
}

func TestMain(m *testing.M) {
	ExpectedRes = make([]int, N)
	for i := range N {
		ExpectedRes[i] = i
	}
	m.Run()
}