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
}