package sink

import (
	"errors"
	"reflect"
	"testing"
	"time"

	"code.justin.tv/ear/kinesis/combiner"
	"code.justin.tv/release/trace/internal/pipeline"
	"golang.org/x/net/context"
)

func TestWriterSink(t *testing.T) {
	ctx := context.Background()
	w := &writer{}
	eh, errs := pipeline.StoreErrors()

	s := WriterSink(ctx, w, eh)

	blobs, want := fakeBlobs(5)
	close(blobs)
	s(blobs)

	for _, err := range *errs {
		t.Errorf("sink err=%q", err)
	}

	have := w.writes
	if !reflect.DeepEqual(have, want) {
		t.Errorf("writer state unexpected")
		t.Logf("have=%v", have)
		t.Logf("want=%v", want)
	}
}

type writer struct {
	writes [][]byte
}

func (w *writer) Write(b []byte) (int, error) {
	w.writes = append(w.writes, b)
	return len(b), nil
}

func TestWriterSinkCancelation(t *testing.T) {
	w := &failWriter{}
	eh, ctx := pipeline.AbortOnError(context.Background())

	s := WriterSink(ctx, w, eh)

	blobs := make(chan *combiner.Blob, 1)
	blobs <- &combiner.Blob{
		Data: []byte("uhoh"),
		Key:  "key",
	}

	done := make(chan struct{})
	go func() {
		s(blobs)
		close(done)
	}()

	timeout := 1 * time.Second
	select {
	case <-time.After(timeout):
		t.Fatal("Timed out waiting for sink to close on error")
	case <-done:
	}
}

type failWriter struct{}

func (w *failWriter) Write(b []byte) (int, error) {
	return 0, errors.New("write failure")
}
