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

     1  // Copyright 2021 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package iofmt_test
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"math/rand"
    11  	"strings"
    12  	"testing"
    13  
    14  	"github.com/Schaudge/grailbase/iofmt"
    15  	"github.com/grailbio/testutil/assert"
    16  )
    17  
    18  // saveWriter saves the calls made to Write for later comparison.
    19  type saveWriter struct {
    20  	writes [][]byte
    21  }
    22  
    23  func (w *saveWriter) Write(p []byte) (int, error) {
    24  	pCopy := make([]byte, len(p))
    25  	copy(pCopy, p)
    26  	w.writes = append(w.writes, pCopy)
    27  	return len(p), nil
    28  }
    29  
    30  // TestLineWriter verifies that a LineWriter calls Write on its underlying
    31  // writer with complete lines.
    32  func TestLineWriter(t *testing.T) {
    33  	for _, c := range []struct {
    34  		name      string
    35  		makeLines func() []string
    36  	}{
    37  		{
    38  			name:      "Empty",
    39  			makeLines: func() []string { return []string{} },
    40  		},
    41  		{
    42  			name: "Numbered",
    43  			makeLines: func() []string {
    44  				const Nlines = 1000
    45  				lines := make([]string, Nlines)
    46  				for i := range lines {
    47  					lines[i] = fmt.Sprintf("line %04d", i)
    48  				}
    49  				return lines
    50  			},
    51  		},
    52  		{
    53  			name: "SomeEmpty",
    54  			makeLines: func() []string {
    55  				const Nlines = 1000
    56  				lines := make([]string, Nlines)
    57  				for i := range lines {
    58  					if rand.Intn(2) == 0 {
    59  						continue
    60  					}
    61  					lines[i] = fmt.Sprintf("line %04d", i)
    62  				}
    63  				return lines
    64  			},
    65  		},
    66  		{
    67  			name: "SomeLong",
    68  			makeLines: func() []string {
    69  				const Nlines = 1000
    70  				lines := make([]string, Nlines)
    71  				for i := range lines {
    72  					var b strings.Builder
    73  					fmt.Fprintf(&b, "line %04d:", i)
    74  					for j := 0; j < rand.Intn(100); j++ {
    75  						b.WriteString(" lorem ipsum")
    76  					}
    77  				}
    78  				return lines
    79  			},
    80  		},
    81  	} {
    82  		t.Run(c.name, func(t *testing.T) {
    83  			lines := c.makeLines()
    84  			// bs is a concatenation of all the lines. We write this to a
    85  			// LineWriter in random segments.
    86  			var bs []byte
    87  			for _, line := range lines {
    88  				bs = append(bs, []byte(fmt.Sprintf("%s\n", line))...)
    89  			}
    90  			s := &saveWriter{}
    91  			w := iofmt.LineWriter(s)
    92  			defer func() {
    93  				assert.Nil(t, w.Close())
    94  			}()
    95  			for len(bs) > 0 {
    96  				// Write in random segments.
    97  				n := rand.Intn(20)
    98  				if len(bs) < n {
    99  					n = len(bs)
   100  				}
   101  				m, err := w.Write(bs[:n])
   102  				assert.Nil(t, err)
   103  				assert.EQ(t, m, n)
   104  				bs = bs[n:]
   105  			}
   106  			want := make([][]byte, len(lines))
   107  			for i, line := range lines {
   108  				want[i] = []byte(fmt.Sprintf("%s\n", line))
   109  			}
   110  			assert.EQ(t, s.writes, want)
   111  		})
   112  	}
   113  }
   114  
   115  // TestLineWriterClose verifies that (*LineWriter).Close writes any remaining
   116  // partial line to the underlying writer.
   117  func TestLineWriterClose(t *testing.T) {
   118  	for _, c := range []struct {
   119  		name string
   120  		bs   []byte
   121  	}{
   122  		{
   123  			name: "Empty",
   124  			bs:   []byte{},
   125  		},
   126  		{
   127  			name: "PartialOnly",
   128  			bs:   []byte("no terminal newline"),
   129  		},
   130  		{
   131  			name: "Partial",
   132  			bs:   []byte("line 0\nline 1\nline 2\nno terminal newline"),
   133  		},
   134  	} {
   135  		t.Run(c.name, func(t *testing.T) {
   136  			var b bytes.Buffer
   137  			w := iofmt.LineWriter(&b)
   138  			_, err := w.Write(c.bs)
   139  			assert.Nil(t, err)
   140  			assert.Nil(t, w.Close())
   141  			assert.EQ(t, b.Bytes(), c.bs)
   142  		})
   143  	}
   144  }