github.com/tcnksm/go@v0.0.0-20141208075154-439b32936367/src/testing/iotest/writer.go (about)

     1  // Copyright 2009 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 iotest
     6  
     7  import "io"
     8  
     9  // TruncateWriter returns a Writer that writes to w
    10  // but stops silently after n bytes.
    11  func TruncateWriter(w io.Writer, n int64) io.Writer {
    12  	return &truncateWriter{w, n}
    13  }
    14  
    15  type truncateWriter struct {
    16  	w io.Writer
    17  	n int64
    18  }
    19  
    20  func (t *truncateWriter) Write(p []byte) (n int, err error) {
    21  	if t.n <= 0 {
    22  		return len(p), nil
    23  	}
    24  	// real write
    25  	n = len(p)
    26  	if int64(n) > t.n {
    27  		n = int(t.n)
    28  	}
    29  	n, err = t.w.Write(p[0:n])
    30  	t.n -= int64(n)
    31  	if err == nil {
    32  		n = len(p)
    33  	}
    34  	return
    35  }