kubo/util/async/forward_test.go

package async

import (
	"testing"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	"github.com/jbenet/go-ipfs/blocks"
)

func TestForwardNThenClose(t *testing.T) {
	const n = 2
	const buf = 2 * n
	in := make(chan *blocks.Block, buf)
	ctx := context.Background()
	out := ForwardN(ctx, in, n)

	for i := 0; i < buf; i++ {
		in <- blocks.NewBlock([]byte(""))
	}

	for i := 0; i < n; i++ {
		_ = <-out
	}

	_, ok := <-out // closed
	if !ok {
		return
	}
	t.Fatal("channel still open after receiving n blocks")
}

func TestCloseInput(t *testing.T) {
	const n = 2
	in := make(chan *blocks.Block, 0)
	ctx := context.Background()
	out := ForwardN(ctx, in, n)

	close(in)
	_, ok := <-out // closed
	if !ok {
		return
	}
	t.Fatal("input channel closed, but output channel not")

}

func TestContextClosedWhenBlockingOnInput(t *testing.T) {
	const n = 1 // but we won't ever send a block
	ctx, cancel := context.WithCancel(context.Background())
	out := ForwardN(ctx, make(chan *blocks.Block), n)

	cancel() // before sending anything
	_, ok := <-out
	if !ok {
		return
	}
	t.Fail()
}