pipeline_test.go

package pipeline

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"regexp"
	"testing"
)

func TestDataFlow(t *testing.T) {
	stages := make([]Stage, 10)
	for i := 0; i < len(stages); i++ {
		stages[i] = testStage{t: t}
	}

	src := &sourceStub{data: stringDataValues(3)}
	sink := new(sinkStub)

	p := NewPipeline(stages...)
	if err := p.Execute(context.TODO(), src, sink); err != nil {
		t.Errorf("Error executing the Pipeline: %v", err)
	}
	if !reflect.DeepEqual(sink.data, src.data) {
		t.Errorf("Data does not match.\nWanted:%v\nGot:%v\n", src.data, sink.data)
	}
}

func TestTaskErrorHandling(t *testing.T) {
	err := errors.New("task error")
	stages := make([]Stage, 10)
	for i := 0; i < len(stages); i++ {
		var sErr error
		if i == 5 {
			sErr = err
		}

		stages[i] = testStage{
			t:   t,
			err: sErr,
		}
	}

	src := &sourceStub{data: stringDataValues(3)}
	sink := new(sinkStub)

	p := NewPipeline(stages...)
	re := regexp.MustCompile("(?s).*task error.*")
	if err := p.Execute(context.TODO(), src, sink); err == nil || !re.MatchString(err.Error()) {
		t.Errorf("Error did not match the expectation: %v", err)
	}
}

type testStage struct {
	id       string
	t        *testing.T
	dropData bool
	err      error
}

func (s testStage) ID() string {
	return s.id
}

func (s testStage) Run(ctx context.Context, sp StageParams) {
	for {
		select {
		case <-ctx.Done():
			return
		case d, ok := <-sp.Input():
			if !ok {
				return
			}
			if s.err != nil {
				sp.Error().Append(s.err)
				return
			}

			if s.dropData {
				continue
			}

			select {
			case <-ctx.Done():
				return
			case sp.Output() <- d:
			}
		}
	}
}

type stringData struct {
	val string
}

func (s *stringData) Clone() Data    { return &stringData{val: s.val} }
func (s *stringData) String() string { return s.val }

func stringDataValues(num int) []Data {
	out := make([]Data, num)

	for i := 0; i < len(out); i++ {
		out[i] = &stringData{val: fmt.Sprint(i)}
	}
	return out
}