github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/helper/gated-writer/writer_test.go (about) 1 package gatedwriter 2 3 import ( 4 "bytes" 5 "io" 6 "strings" 7 "sync" 8 "testing" 9 "time" 10 11 "github.com/stretchr/testify/require" 12 ) 13 14 func TestWriter_impl(t *testing.T) { 15 var _ io.Writer = new(Writer) 16 } 17 18 type slowTestWriter struct { 19 buf *bytes.Buffer 20 called chan struct{} 21 callCount int 22 } 23 24 func (w *slowTestWriter) Write(p []byte) (int, error) { 25 if w.callCount == 0 { 26 defer close(w.called) 27 } 28 29 w.callCount++ 30 time.Sleep(time.Millisecond) 31 32 return w.buf.Write(p) 33 } 34 35 func TestWriter_WithSlowWriter(t *testing.T) { 36 buf := new(bytes.Buffer) 37 called := make(chan struct{}) 38 39 w := &slowTestWriter{ 40 buf: buf, 41 called: called, 42 } 43 44 writer := &Writer{Writer: w} 45 46 writer.Write([]byte("foo\n")) 47 writer.Write([]byte("bar\n")) 48 writer.Write([]byte("baz\n")) 49 50 flushed := make(chan struct{}) 51 52 go func() { 53 writer.Flush() 54 close(flushed) 55 }() 56 57 // wait for the flush to call Write on slowTestWriter 58 <-called 59 60 // write to the now-flushing writer, which is no longer buffering 61 writer.Write([]byte("quux\n")) 62 63 // wait for the flush to finish to assert 64 <-flushed 65 66 require.Equal(t, "foo\nbar\nbaz\nquux\n", buf.String()) 67 } 68 69 func TestWriter(t *testing.T) { 70 buf := new(bytes.Buffer) 71 w := &Writer{Writer: buf} 72 w.Write([]byte("foo\n")) 73 w.Write([]byte("bar\n")) 74 75 if buf.String() != "" { 76 t.Fatalf("bad: %s", buf.String()) 77 } 78 79 w.Flush() 80 81 if buf.String() != "foo\nbar\n" { 82 t.Fatalf("bad: %s", buf.String()) 83 } 84 85 w.Write([]byte("baz\n")) 86 87 if buf.String() != "foo\nbar\nbaz\n" { 88 t.Fatalf("bad: %s", buf.String()) 89 } 90 } 91 92 func TestWriter_WithMultipleWriters(t *testing.T) { 93 buf := new(bytes.Buffer) 94 95 writer := &Writer{Writer: buf} 96 97 strs := []string{ 98 "foo\n", 99 "bar\n", 100 "baz\n", 101 "quux\n", 102 } 103 104 waitCh := make(chan struct{}) 105 106 wg := &sync.WaitGroup{} 107 108 for _, str := range strs { 109 str := str 110 wg.Add(1) 111 go func() { 112 defer wg.Done() 113 <-waitCh 114 writer.Write([]byte(str)) 115 }() 116 } 117 118 // synchronize calls to Write() as closely as possible 119 close(waitCh) 120 121 wg.Wait() 122 123 writer.Flush() 124 125 require.Equal(t, strings.Count(buf.String(), "\n"), len(strs)) 126 }