path: root/runtime/task_test.go

package runtime

import (
	"context"
	"errors"
	"testing"

	"goflink/core"
	"goflink/core/operator"
	"goflink/state"
	"goflink/transport"
)

// map(x*10) -> filter(>20), end to end.
func TestPipeline(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	source := make(chan core.StreamRecord[int])
	mid := make(chan core.StreamRecord[int])
	sink := make(chan core.StreamRecord[int])

	mapped := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(mid),
		operator.NewMapOperator(func(in int) (int, error) { return in * 10, nil }),
	)
	filtered := RunTask(ctx,
		transport.NewLocalReceiveChannel(mid),
		transport.NewLocalEmitChannel(sink),
		operator.NewFilterOperator(func(in int) bool { return in > 20 }),
	)

	go func() {
		defer close(source)
		for i := 1; i <= 5; i++ {
			source <- core.NewDataRecord("k", i, int64(i))
		}
	}()

	var got []int
	for record := range sink {
		got = append(got, record.Value)
	}

	want := []int{30, 40, 50}
	if len(got) != len(want) {
		t.Fatalf("got %v, want %v", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("got %v, want %v", got, want)
		}
	}

	for _, errCh := range []<-chan error{mapped, filtered} {
		if err := <-errCh; err != nil {
			t.Fatalf("task failed: %v", err)
		}
	}
}

// wordCounter is the canonical stateful user func: read the old count for the
// current key, +1, store it back.
type wordCounter struct {
	count core.ValueState[any]
}

func (w *wordCounter) Open(ctx context.Context) error {
	be, ok := core.ExtractStateBackend(ctx)
	if !ok {
		return errors.New("no state backend in context")
	}
	s, err := be.GetValueState(ctx, "word-count")
	if err != nil {
		return err
	}
	w.count = s
	return nil
}

func (w *wordCounter) Map(ctx context.Context, in string) (int, error) {
	v, err := w.count.Value(ctx)
	if err != nil {
		return 0, err
	}
	n, _ := v.(int) // unseen key -> nil -> 0
	n++
	_, err = w.count.Update(ctx, n)
	return n, err
}

func (w *wordCounter) Close(ctx context.Context) error { return nil }

// keyBy -> statefulMap: state must be per key and survive across records.
func TestStatefulMapCountsPerKey(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))

	source := make(chan core.StreamRecord[string])
	mid := make(chan core.StreamRecord[string])
	sink := make(chan core.StreamRecord[int])

	keyed := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(mid),
		operator.NewKeyByOperator(func(in string) string { return in }),
	)
	counted := RunTask(ctx,
		transport.NewLocalReceiveChannel(mid),
		transport.NewLocalEmitChannel(sink),
		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
	)

	go func() {
		defer close(source)
		for _, name := range []string{"Sang", "Tran", "Sang", "Sang"} {
			source <- core.NewDataRecord("", name, 0)
		}
	}()

	var got []int
	for record := range sink {
		got = append(got, record.Value)
	}

	want := []int{1, 1, 2, 3}
	if len(got) != len(want) {
		t.Fatalf("got %v, want %v", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("got %v, want %v", got, want)
		}
	}

	for _, errCh := range []<-chan error{keyed, counted} {
		if err := <-errCh; err != nil {
			t.Fatalf("task failed: %v", err)
		}
	}
}

// An unkeyed record must fail loudly instead of silently sharing one state slot.
func TestStatefulMapRejectsUnkeyed(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctx = core.InjectStateBackend(ctx, state.NewMemoryStateBackend(core.DefaultMaxParallelism))

	source := make(chan core.StreamRecord[string], 1)
	sink := make(chan core.StreamRecord[int], 1)

	done := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(sink),
		operator.NewStatefulMapOperator[string, int](&wordCounter{}),
	)

	source <- core.NewDataRecord("", "Sang", 0)

	if err := <-done; err == nil {
		t.Fatal("unkeyed record was accepted, want an error")
	}
}

// A cancelled ctx must unblock a task parked in Receive.
func TestCancelStopsTask(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	source := make(chan core.StreamRecord[int])
	sink := make(chan core.StreamRecord[int], 1)

	done := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(sink),
		operator.NewMapOperator(func(in int) (int, error) { return in, nil }),
	)

	cancel()

	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("got %v, want context.Canceled", err)
	}
}

// A user function error must reach the caller with its cause intact.
func TestProcessErrorPropagates(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	boom := errors.New("boom")
	source := make(chan core.StreamRecord[int], 1)
	sink := make(chan core.StreamRecord[int], 1)

	done := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(sink),
		operator.NewMapOperator(func(in int) (int, error) { return 0, boom }),
	)

	source <- core.NewDataRecord("k", 1, 0)

	if err := <-done; !errors.Is(err, boom) {
		t.Fatalf("got %v, want %v", err, boom)
	}
}

// closeProbe records the ctx.Err() its Close saw.
type closeProbe struct {
	closeCtxErr error
}

func (o *closeProbe) Open(ctx context.Context) error { return nil }

func (o *closeProbe) Process(ctx context.Context, r core.StreamRecord[int], out operator.Emitter[int]) error {
	return out.Emit(ctx, r)
}

func (o *closeProbe) Close(ctx context.Context) error {
	o.closeCtxErr = ctx.Err()
	return nil
}

// Close must get a live ctx even when the task was killed by cancellation,
// otherwise an operator that flushes on shutdown loses whatever it buffered.
func TestCloseGetsLiveContext(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	source := make(chan core.StreamRecord[int])
	sink := make(chan core.StreamRecord[int], 1)

	probe := &closeProbe{}
	done := RunTask(ctx,
		transport.NewLocalReceiveChannel(source),
		transport.NewLocalEmitChannel(sink),
		probe,
	)

	cancel()
	<-done // Close has already run by the time this channel closes

	if probe.closeCtxErr != nil {
		t.Fatalf("Close got a dead ctx (%v), so it cannot flush", probe.closeCtxErr)
	}
}