github.com/khulnasoft/cli@v0.0.0-20240402070845-01bcad7beefa/internal/test/writer.go (about)

     1  package test
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  // WriterWithHook is an io.Writer that calls a hook function
     8  // after every write.
     9  // This is useful in testing to wait for a write to complete,
    10  // or to check what was written.
    11  // To create a WriterWithHook use the NewWriterWithHook function.
    12  type WriterWithHook struct {
    13  	actualWriter io.Writer
    14  	hook         func([]byte)
    15  }
    16  
    17  // Write writes p to the actual writer and then calls the hook function.
    18  func (w *WriterWithHook) Write(p []byte) (n int, err error) {
    19  	defer w.hook(p)
    20  	return w.actualWriter.Write(p)
    21  }
    22  
    23  var _ io.Writer = (*WriterWithHook)(nil)
    24  
    25  // NewWriterWithHook returns a new WriterWithHook that still writes to the actualWriter
    26  // but also calls the hook function after every write.
    27  // The hook function is useful for testing, or waiting for a write to complete.
    28  func NewWriterWithHook(actualWriter io.Writer, hook func([]byte)) *WriterWithHook {
    29  	return &WriterWithHook{actualWriter: actualWriter, hook: hook}
    30  }