multi_test.go

     1  // Copyright 2010 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package io_test
     6  
     7  import (
     8  	"bytes"
     9  	"crypto/sha1"
    10  	"errors"
    11  	"fmt"
    12  	. "io"
    13  	"runtime"
    14  	"strings"
    15  	"testing"
    16  	"time"
    17  )
    18  
    19  func TestMultiReader(t *testing.T) {
    20  	var mr Reader
    21  	var buf []byte
    22  	nread := 0
    23  	withFooBar := func(tests func()) {
    24  		r1 := strings.NewReader("foo ")
    25  		r2 := strings.NewReader("")
    26  		r3 := strings.NewReader("bar")
    27  		mr = MultiReader(r1, r2, r3)
    28  		buf = make([]byte, 20)
    29  		tests()
    30  	}
    31  	expectRead := func(size int, expected string, eerr error) {
    32  		nread++
    33  		n, gerr := mr.Read(buf[0:size])
    34  		if n != len(expected) {
    35  			t.Errorf("#%d, expected %d bytes; got %d",
    36  				nread, len(expected), n)
    37  		}
    38  		got := string(buf[0:n])
    39  		if got != expected {
    40  			t.Errorf("#%d, expected %q; got %q",
    41  				nread, expected, got)
    42  		}
    43  		if gerr != eerr {
    44  			t.Errorf("#%d, expected error %v; got %v",
    45  				nread, eerr, gerr)
    46  		}
    47  		buf = buf[n:]
    48  	}
    49  	withFooBar(func() {
    50  		expectRead(2, "fo", nil)
    51  		expectRead(5, "o ", nil)
    52  		expectRead(5, "bar", nil)
    53  		expectRead(5, "", EOF)
    54  	})
    55  	withFooBar(func() {
    56  		expectRead(4, "foo ", nil)
    57  		expectRead(1, "b", nil)
    58  		expectRead(3, "ar", nil)
    59  		expectRead(1, "", EOF)
    60  	})
    61  	withFooBar(func() {
    62  		expectRead(5, "foo ", nil)
    63  	})
    64  }
    65  
    66  func TestMultiReaderAsWriterTo(t *testing.T) {
    67  	mr := MultiReader(
    68  		strings.NewReader("foo "),
    69  		MultiReader( // Tickle the buffer reusing codepath
    70  			strings.NewReader(""),
    71  			strings.NewReader("bar"),
    72  		),
    73  	)
    74  	mrAsWriterTo, ok := mr.(WriterTo)
    75  	if !ok {
    76  		t.Fatalf("expected cast to WriterTo to succeed")
    77  	}
    78  	sink := &strings.Builder{}
    79  	n, err := mrAsWriterTo.WriteTo(sink)
    80  	if err != nil {
    81  		t.Fatalf("expected no error; got %v", err)
    82  	}
    83  	if n != 7 {
    84  		t.Errorf("expected read 7 bytes; got %d", n)
    85  	}
    86  	if result := sink.String(); result != "foo bar" {
    87  		t.Errorf(`expected "foo bar"; got %q`, result)
    88  	}
    89  }
    90  
    91  func TestMultiWriter(t *testing.T) {
    92  	sink := new(bytes.Buffer)
    93  	// Hide bytes.Buffer's WriteString method:
    94  	testMultiWriter(t, struct {
    95  		Writer
    96  		fmt.Stringer
    97  	}{sink, sink})
    98  }
    99  
   100  func TestMultiWriter_String(t *testing.T) {
   101  	testMultiWriter(t, new(bytes.Buffer))
   102  }
   103  
   104  // Test that a multiWriter.WriteString calls results in at most 1 allocation,
   105  // even if multiple targets don't support WriteString.
   106  func TestMultiWriter_WriteStringSingleAlloc(t *testing.T) {
   107  	var sink1, sink2 bytes.Buffer
   108  	type simpleWriter struct { // hide bytes.Buffer's WriteString
   109  		Writer
   110  	}
   111  	mw := MultiWriter(simpleWriter{&sink1}, simpleWriter{&sink2})
   112  	allocs := int(testing.AllocsPerRun(1000, func() {
   113  		WriteString(mw, "foo")
   114  	}))
   115  	if allocs != 1 {
   116  		t.Errorf("num allocations = %d; want 1", allocs)
   117  	}
   118  }
   119  
   120  type writeStringChecker struct{ called bool }
   121  
   122  func (c *writeStringChecker) WriteString(s string) (n int, err error) {
   123  	c.called = true
   124  	return len(s), nil
   125  }
   126  
   127  func (c *writeStringChecker) Write(p []byte) (n int, err error) {
   128  	return len(p), nil
   129  }
   130  
   131  func TestMultiWriter_StringCheckCall(t *testing.T) {
   132  	var c writeStringChecker
   133  	mw := MultiWriter(&c)
   134  	WriteString(mw, "foo")
   135  	if !c.called {
   136  		t.Error("did not see WriteString call to writeStringChecker")
   137  	}
   138  }
   139  
   140  func testMultiWriter(t *testing.T, sink interface {
   141  	Writer
   142  	fmt.Stringer
   143  }) {
   144  	sha1 := sha1.New()
   145  	mw := MultiWriter(sha1, sink)
   146  
   147  	sourceString := "My input text."
   148  	source := strings.NewReader(sourceString)
   149  	written, err := Copy(mw, source)
   150  
   151  	if written != int64(len(sourceString)) {
   152  		t.Errorf("short write of %d, not %d", written, len(sourceString))
   153  	}
   154  
   155  	if err != nil {
   156  		t.Errorf("unexpected error: %v", err)
   157  	}
   158  
   159  	sha1hex := fmt.Sprintf("%x", sha1.Sum(nil))
   160  	if sha1hex != "01cb303fa8c30a64123067c5aa6284ba7ec2d31b" {
   161  		t.Error("incorrect sha1 value")
   162  	}
   163  
   164  	if sink.String() != sourceString {
   165  		t.Errorf("expected %q; got %q", sourceString, sink.String())
   166  	}
   167  }
   168  
   169  // writerFunc is a Writer implemented by the underlying func.
   170  type writerFunc func(p []byte) (int, error)
   171  
   172  func (f writerFunc) Write(p []byte) (int, error) {
   173  	return f(p)
   174  }
   175  
   176  // Test that MultiWriter properly flattens chained multiWriters.
   177  func TestMultiWriterSingleChainFlatten(t *testing.T) {
   178  	pc := make([]uintptr, 1000) // 1000 should fit the full stack
   179  	n := runtime.Callers(0, pc)
   180  	var myDepth = callDepth(pc[:n])
   181  	var writeDepth int // will contain the depth from which writerFunc.Writer was called
   182  	var w Writer = MultiWriter(writerFunc(func(p []byte) (int, error) {
   183  		n := runtime.Callers(1, pc)
   184  		writeDepth += callDepth(pc[:n])
   185  		return 0, nil
   186  	}))
   187  
   188  	mw := w
   189  	// chain a bunch of multiWriters
   190  	for i := 0; i < 100; i++ {
   191  		mw = MultiWriter(w)
   192  	}
   193  
   194  	mw = MultiWriter(w, mw, w, mw)
   195  	mw.Write(nil) // don't care about errors, just want to check the call-depth for Write
   196  
   197  	if writeDepth != 4*(myDepth+2) { // 2 should be multiWriter.Write and writerFunc.Write
   198  		t.Errorf("multiWriter did not flatten chained multiWriters: expected writeDepth %d, got %d",
   199  			4*(myDepth+2), writeDepth)
   200  	}
   201  }
   202  
   203  func TestMultiWriterError(t *testing.T) {
   204  	f1 := writerFunc(func(p []byte) (int, error) {
   205  		return len(p) / 2, ErrShortWrite
   206  	})
   207  	f2 := writerFunc(func(p []byte) (int, error) {
   208  		t.Errorf("MultiWriter called f2.Write")
   209  		return len(p), nil
   210  	})
   211  	w := MultiWriter(f1, f2)
   212  	n, err := w.Write(make([]byte, 100))
   213  	if n != 50 || err != ErrShortWrite {
   214  		t.Errorf("Write = %d, %v, want 50, ErrShortWrite", n, err)
   215  	}
   216  }
   217  
   218  // Test that MultiReader copies the input slice and is insulated from future modification.
   219  func TestMultiReaderCopy(t *testing.T) {
   220  	slice := []Reader{strings.NewReader("hello world")}
   221  	r := MultiReader(slice...)
   222  	slice[0] = nil
   223  	data, err := ReadAll(r)
   224  	if err != nil || string(data) != "hello world" {
   225  		t.Errorf("ReadAll() = %q, %v, want %q, nil", data, err, "hello world")
   226  	}
   227  }
   228  
   229  // Test that MultiWriter copies the input slice and is insulated from future modification.
   230  func TestMultiWriterCopy(t *testing.T) {
   231  	var buf strings.Builder
   232  	slice := []Writer{&buf}
   233  	w := MultiWriter(slice...)
   234  	slice[0] = nil
   235  	n, err := w.Write([]byte("hello world"))
   236  	if err != nil || n != 11 {
   237  		t.Errorf("Write(`hello world`) = %d, %v, want 11, nil", n, err)
   238  	}
   239  	if buf.String() != "hello world" {
   240  		t.Errorf("buf.String() = %q, want %q", buf.String(), "hello world")
   241  	}
   242  }
   243  
   244  // readerFunc is a Reader implemented by the underlying func.
   245  type readerFunc func(p []byte) (int, error)
   246  
   247  func (f readerFunc) Read(p []byte) (int, error) {
   248  	return f(p)
   249  }
   250  
   251  // callDepth returns the logical call depth for the given PCs.
   252  func callDepth(callers []uintptr) (depth int) {
   253  	frames := runtime.CallersFrames(callers)
   254  	more := true
   255  	for more {
   256  		_, more = frames.Next()
   257  		depth++
   258  	}
   259  	return
   260  }
   261  
   262  // Test that MultiReader properly flattens chained multiReaders when Read is called
   263  func TestMultiReaderFlatten(t *testing.T) {
   264  	pc := make([]uintptr, 1000) // 1000 should fit the full stack
   265  	n := runtime.Callers(0, pc)
   266  	var myDepth = callDepth(pc[:n])
   267  	var readDepth int // will contain the depth from which fakeReader.Read was called
   268  	var r Reader = MultiReader(readerFunc(func(p []byte) (int, error) {
   269  		n := runtime.Callers(1, pc)
   270  		readDepth = callDepth(pc[:n])
   271  		return 0, errors.New("irrelevant")
   272  	}))
   273  
   274  	// chain a bunch of multiReaders
   275  	for i := 0; i < 100; i++ {
   276  		r = MultiReader(r)
   277  	}
   278  
   279  	r.Read(nil) // don't care about errors, just want to check the call-depth for Read
   280  
   281  	if readDepth != myDepth+2 { // 2 should be multiReader.Read and fakeReader.Read
   282  		t.Errorf("multiReader did not flatten chained multiReaders: expected readDepth %d, got %d",
   283  			myDepth+2, readDepth)
   284  	}
   285  }
   286  
   287  // byteAndEOFReader is a Reader which reads one byte (the underlying
   288  // byte) and EOF at once in its Read call.
   289  type byteAndEOFReader byte
   290  
   291  func (b byteAndEOFReader) Read(p []byte) (n int, err error) {
   292  	if len(p) == 0 {
   293  		// Read(0 bytes) is useless. We expect no such useless
   294  		// calls in this test.
   295  		panic("unexpected call")
   296  	}
   297  	p[0] = byte(b)
   298  	return 1, EOF
   299  }
   300  
   301  // This used to yield bytes forever; issue 16795.
   302  func TestMultiReaderSingleByteWithEOF(t *testing.T) {
   303  	got, err := ReadAll(LimitReader(MultiReader(byteAndEOFReader('a'), byteAndEOFReader('b')), 10))
   304  	if err != nil {
   305  		t.Fatal(err)
   306  	}
   307  	const want = "ab"
   308  	if string(got) != want {
   309  		t.Errorf("got %q; want %q", got, want)
   310  	}
   311  }
   312  
   313  // Test that a reader returning (n, EOF) at the end of a MultiReader
   314  // chain continues to return EOF on its final read, rather than
   315  // yielding a (0, EOF).
   316  func TestMultiReaderFinalEOF(t *testing.T) {
   317  	r := MultiReader(bytes.NewReader(nil), byteAndEOFReader('a'))
   318  	buf := make([]byte, 2)
   319  	n, err := r.Read(buf)
   320  	if n != 1 || err != EOF {
   321  		t.Errorf("got %v, %v; want 1, EOF", n, err)
   322  	}
   323  }
   324  
   325  func TestMultiReaderFreesExhaustedReaders(t *testing.T) {
   326  	var mr Reader
   327  	closed := make(chan struct{})
   328  	// The closure ensures that we don't have a live reference to buf1
   329  	// on our stack after MultiReader is inlined (Issue 18819).  This
   330  	// is a work around for a limitation in liveness analysis.
   331  	func() {
   332  		buf1 := bytes.NewReader([]byte("foo"))
   333  		buf2 := bytes.NewReader([]byte("bar"))
   334  		mr = MultiReader(buf1, buf2)
   335  		runtime.SetFinalizer(buf1, func(*bytes.Reader) {
   336  			close(closed)
   337  		})
   338  	}()
   339  
   340  	buf := make([]byte, 4)
   341  	if n, err := ReadFull(mr, buf); err != nil || string(buf) != "foob" {
   342  		t.Fatalf(`ReadFull = %d (%q), %v; want 3, "foo", nil`, n, buf[:n], err)
   343  	}
   344  
   345  	runtime.GC()
   346  	select {
   347  	case <-closed:
   348  	case <-time.After(5 * time.Second):
   349  		t.Fatal("timeout waiting for collection of buf1")
   350  	}
   351  
   352  	if n, err := ReadFull(mr, buf[:2]); err != nil || string(buf[:2]) != "ar" {
   353  		t.Fatalf(`ReadFull = %d (%q), %v; want 2, "ar", nil`, n, buf[:n], err)
   354  	}
   355  }
   356  
   357  func TestInterleavedMultiReader(t *testing.T) {
   358  	r1 := strings.NewReader("123")
   359  	r2 := strings.NewReader("45678")
   360  
   361  	mr1 := MultiReader(r1, r2)
   362  	mr2 := MultiReader(mr1)
   363  
   364  	buf := make([]byte, 4)
   365  
   366  	// Have mr2 use mr1's []Readers.
   367  	// Consume r1 (and clear it for GC to handle) and consume part of r2.
   368  	n, err := ReadFull(mr2, buf)
   369  	if got := string(buf[:n]); got != "1234" || err != nil {
   370  		t.Errorf(`ReadFull(mr2) = (%q, %v), want ("1234", nil)`, got, err)
   371  	}
   372  
   373  	// Consume the rest of r2 via mr1.
   374  	// This should not panic even though mr2 cleared r1.
   375  	n, err = ReadFull(mr1, buf)
   376  	if got := string(buf[:n]); got != "5678" || err != nil {
   377  		t.Errorf(`ReadFull(mr1) = (%q, %v), want ("5678", nil)`, got, err)
   378  	}
   379  }
   380
View as plain text