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  }