github.com/mutagen-io/mutagen@v0.18.0-rc1/pkg/stream/preemptable_writer_test.go (about)

     1  package stream
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"math/rand"
     7  	"testing"
     8  )
     9  
    10  // TestPreemptableWriter tests preemptableWriter.
    11  func TestPreemptableWriter(t *testing.T) {
    12  	// Set up parameters.
    13  	const (
    14  		checkInterval  = 3
    15  		copyBufferSize = 1024
    16  	)
    17  
    18  	// Create a limited input stream. Ensure that it's large enough to trigger a
    19  	// preemption check when copied using our copy buffer.
    20  	source := &io.LimitedReader{
    21  		R: rand.New(rand.NewSource(0)),
    22  		N: (checkInterval + 1) * copyBufferSize,
    23  	}
    24  
    25  	// Create a copy buffer.
    26  	copyBuffer := make([]byte, copyBufferSize)
    27  
    28  	// Create a closed channel so that the write is already preempted.
    29  	cancelled := make(chan struct{})
    30  	close(cancelled)
    31  
    32  	// Create a preemptable writer.
    33  	destinationBuffer := &bytes.Buffer{}
    34  	destination := NewPreemptableWriter(destinationBuffer, cancelled, 3)
    35  
    36  	// Perform a copy.
    37  	n, err := io.CopyBuffer(destination, source, copyBuffer)
    38  
    39  	// Check the copy error.
    40  	if err == nil {
    41  		t.Error("copy not preempted")
    42  	} else if err != ErrWritePreempted {
    43  		t.Fatal("unexpected copy error:", err)
    44  	}
    45  
    46  	// Check the copy sizes.
    47  	if n != int64(destinationBuffer.Len()) {
    48  		t.Error("preemptable writer did not write reported length")
    49  	}
    50  	if n != checkInterval*copyBufferSize {
    51  		t.Error("unexpected number of bytes written between preemption checks")
    52  	}
    53  }