package sink

import (
	"fmt"
	"math/rand"
	"sync/atomic"
	"testing"

	"code.justin.tv/ear/kinesis/combiner"
)

func TestSampleSink(t *testing.T) {
	rand.Seed(0)
	var nRecv int64 = 0
	nSent := 100
	rate := 0.1
	nWant := 9

	counter := counterSink(&nRecv)
	sampled := Sample(rate, counter)

	blobs, _ := fakeBlobs(nSent)
	close(blobs)
	sampled(blobs)

	if atomic.LoadInt64(&nRecv) != int64(nWant) {
		t.Errorf("unexpected # after sampling have=%d want=%d", nRecv, nWant)
	}
}

// a sink that just counts blobs
func counterSink(counter *int64) Sink {
	return func(data <-chan *combiner.Blob) {
		for _ = range data {
			atomic.AddInt64(counter, 1)
		}
	}
}

func fakeBlobs(n int) (ch chan *combiner.Blob, want [][]byte) {
	ch = make(chan *combiner.Blob, n)
	want = make([][]byte, n)
	for i := 0; i < n; i++ {
		want[i] = []byte(fmt.Sprintf("data-%d", i))
		ch <- &combiner.Blob{
			Data: want[i],
			Key:  fmt.Sprintf("key-%d", i),
		}
	}
	return ch, want
}
