github.com/hernad/nomad@v1.6.112/helper/gated-writer/writer_test.go (about)

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