34 lines
442 B
Go
34 lines
442 B
Go
|
package testutil
|
||
|
|
|
||
|
|
import (
|
||
|
|
"io"
|
||
|
|
"os"
|
||
|
|
)
|
||
|
|
|
||
|
|
func NewLimitWriter(w io.Writer, limit int) *LimitWriter {
|
||
|
|
return &LimitWriter{
|
||
|
|
w: w,
|
||
|
|
limit: limit,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type LimitWriter struct {
|
||
|
|
w io.Writer
|
||
|
|
limit int
|
||
|
|
written int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (lw *LimitWriter) Write(buf []byte) (int, error) {
|
||
|
|
n, err := lw.w.Write(buf)
|
||
|
|
if err != nil {
|
||
|
|
return n, err
|
||
|
|
}
|
||
|
|
|
||
|
|
lw.written += n
|
||
|
|
if lw.written > lw.limit {
|
||
|
|
return n, os.ErrClosed
|
||
|
|
}
|
||
|
|
|
||
|
|
return n, nil
|
||
|
|
}
|