github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/limitbuf/limitbuf_test.go (about)

     1  package limitbuf_test
     2  
     3  import (
     4  	"bytes"
     5  	"testing"
     6  
     7  	"github.com/Schaudge/grailbase/limitbuf"
     8  	"github.com/Schaudge/grailbase/log"
     9  	"github.com/grailbio/testutil/expect"
    10  )
    11  
    12  func TestLogger(t *testing.T) {
    13  	l := limitbuf.NewLogger(10)
    14  	l.Write([]byte("blah"))
    15  	expect.EQ(t, l.String(), "blah")
    16  	l.Write([]byte("abcdefgh"))
    17  	expect.EQ(t, l.String(), "blahabcdef(truncated 2 bytes)")
    18  	expect.EQ(t, l.String(), "blahabcdef(truncated 2 bytes)")
    19  }
    20  
    21  func TestLoggerExtremeTruncation(t *testing.T) {
    22  	oldOutputter := log.GetOutputter()
    23  	t.Cleanup(func() { log.SetOutputter(oldOutputter) })
    24  	var outputter testOutputter
    25  	log.SetOutputter(&outputter)
    26  
    27  	logger := limitbuf.NewLogger(2, limitbuf.LogIfTruncatingMaxMultiple(3))
    28  	_, err := logger.Write([]byte("abcdefg"))
    29  	expect.NoError(t, err)
    30  
    31  	expect.EQ(t, logger.String(), "ab(truncated 5 bytes)")
    32  	expect.HasSubstr(t, outputter.String(), "extreme truncation")
    33  }
    34  
    35  type testOutputter struct{ bytes.Buffer }
    36  
    37  func (o *testOutputter) Level() log.Level {
    38  	return log.Error
    39  }
    40  func (o *testOutputter) Output(_ int, _ log.Level, s string) error {
    41  	_, err := o.Buffer.WriteString(s)
    42  	return err
    43  }