github.com/blend/go-sdk@v1.20220411.3/logger/interlocked_writer_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package logger
     9  
    10  import (
    11  	"bytes"
    12  	"io"
    13  	"testing"
    14  
    15  	"github.com/blend/go-sdk/assert"
    16  )
    17  
    18  var (
    19  	_ io.WriteCloser = (*mockWriter)(nil)
    20  )
    21  
    22  // mockWriter is a stub for a io.WriteCloser.
    23  type mockWriter struct {
    24  	WriteHandler func([]byte) (int, error)
    25  	CloseHandler func() error
    26  }
    27  
    28  // Write implements io.Writer
    29  func (mw mockWriter) Write(data []byte) (int, error) {
    30  	if mw.WriteHandler != nil {
    31  		return mw.WriteHandler(data)
    32  	}
    33  	return 0, nil
    34  }
    35  
    36  // Close implements io.Closer.
    37  func (mw mockWriter) Close() error {
    38  	if mw.CloseHandler != nil {
    39  		return mw.CloseHandler()
    40  	}
    41  	return nil
    42  }
    43  
    44  func TestInterlockedWriter(t *testing.T) {
    45  	assert := assert.New(t)
    46  
    47  	buf := new(bytes.Buffer)
    48  	var didWrite, didClose bool
    49  	mw := mockWriter{
    50  		WriteHandler: func(data []byte) (int, error) {
    51  			defer func() {
    52  				didWrite = true
    53  			}()
    54  			return buf.Write(data)
    55  		},
    56  		CloseHandler: func() error {
    57  			defer func() {
    58  				didClose = true
    59  			}()
    60  			return nil
    61  		},
    62  	}
    63  
    64  	iw := NewInterlockedWriter(mw)
    65  	data := []byte("this is a test")
    66  	written, err := iw.Write(data)
    67  	assert.Nil(err)
    68  	assert.Equal(len(data), written)
    69  	assert.True(didWrite)
    70  	assert.Nil(iw.Close())
    71  	assert.True(didClose)
    72  
    73  	assert.Equal("this is a test", buf.String())
    74  }