github.com/quic-go/quic-go@v0.44.0/qlog/writer_test.go (about)

     1  package qlog
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"io"
     7  	"log"
     8  	"os"
     9  
    10  	"github.com/quic-go/quic-go/internal/protocol"
    11  
    12  	. "github.com/onsi/ginkgo/v2"
    13  	. "github.com/onsi/gomega"
    14  )
    15  
    16  type limitedWriter struct {
    17  	io.WriteCloser
    18  	N       int
    19  	written int
    20  }
    21  
    22  func (w *limitedWriter) Write(p []byte) (int, error) {
    23  	if w.written+len(p) > w.N {
    24  		return 0, errors.New("writer full")
    25  	}
    26  	n, err := w.WriteCloser.Write(p)
    27  	w.written += n
    28  	return n, err
    29  }
    30  
    31  var _ = Describe("Writing", func() {
    32  	It("stops writing when encountering an error", func() {
    33  		buf := &bytes.Buffer{}
    34  		t := NewConnectionTracer(
    35  			&limitedWriter{WriteCloser: nopWriteCloser(buf), N: 250},
    36  			protocol.PerspectiveServer,
    37  			protocol.ParseConnectionID([]byte{0xde, 0xad, 0xbe, 0xef}),
    38  		)
    39  		for i := uint32(0); i < 1000; i++ {
    40  			t.UpdatedPTOCount(i)
    41  		}
    42  
    43  		b := &bytes.Buffer{}
    44  		log.SetOutput(b)
    45  		defer log.SetOutput(os.Stdout)
    46  		t.Close()
    47  		Expect(b.String()).To(ContainSubstring("writer full"))
    48  	})
    49  })