github.com/pachyderm/pachyderm@v1.13.4/src/server/worker/logs/testing.go (about)

     1  package logs
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/pachyderm/pachyderm/src/client/pfs"
     9  	"github.com/pachyderm/pachyderm/src/client/pkg/errors"
    10  	"github.com/pachyderm/pachyderm/src/server/worker/common"
    11  )
    12  
    13  // MockLogger is an implementation of the TaggedLogger interface for use in
    14  // tests.  Loggers are often passed to callbacks, so you can check that the
    15  // logger has been configured with the right tags in these cases.  In addition,
    16  // you can set the Writer field so that log statements go directly to stdout
    17  // (or some other location) for debugging purposes.
    18  type MockLogger struct {
    19  	// These fields are exposed so that tests can fuck around with them or make assertions
    20  	Writer   io.Writer
    21  	Job      string
    22  	Data     []*common.Input
    23  	UserCode bool
    24  }
    25  
    26  // Not used - forces a compile-time error in this file if MockLogger does not
    27  // implement TaggedLogger
    28  var _ TaggedLogger = &MockLogger{}
    29  
    30  // NewMockLogger constructs a MockLogger object for use by tests.
    31  func NewMockLogger() *MockLogger {
    32  	return &MockLogger{}
    33  }
    34  
    35  // Write fulfills the io.Writer interface for TaggedLogger, and will optionally
    36  // write to the configured ml.Writer, otherwise it pretends that it succeeded.
    37  func (ml *MockLogger) Write(p []byte) (_ int, retErr error) {
    38  	if ml.Writer != nil {
    39  		return ml.Writer.Write(p)
    40  	}
    41  	return len(p), nil
    42  }
    43  
    44  // Logf optionally logs a statement using string formatting
    45  func (ml *MockLogger) Logf(formatString string, args ...interface{}) {
    46  	if ml.Writer != nil {
    47  		params := []interface{}{time.Now().Format(time.StampMilli), ml.Job, ml.Data, ml.UserCode}
    48  		params = append(params, args...)
    49  		str := fmt.Sprintf("LOGF %s (%v, %v, %v): "+formatString+"\n", params...)
    50  		ml.Writer.Write([]byte(str))
    51  	}
    52  }
    53  
    54  // Errf optionally logs an error statement using string formatting
    55  func (ml *MockLogger) Errf(formatString string, args ...interface{}) {
    56  	if ml.Writer != nil {
    57  		params := []interface{}{time.Now().Format(time.StampMilli), ml.Job, ml.Data, ml.UserCode}
    58  		params = append(params, args...)
    59  		str := fmt.Sprintf("ERRF %s (%v, %v, %v): "+formatString+"\n", params...)
    60  		ml.Writer.Write([]byte(str))
    61  	}
    62  }
    63  
    64  // LogStep will log before and after the given callback function runs, using
    65  // the name provided
    66  func (ml *MockLogger) LogStep(name string, cb func() error) (retErr error) {
    67  	ml.Logf("started %v", name)
    68  	defer func() {
    69  		if retErr != nil {
    70  			retErr = errors.EnsureStack(retErr)
    71  			ml.Logf("errored %v: %v", name, retErr)
    72  		} else {
    73  			ml.Logf("finished %v", name)
    74  		}
    75  	}()
    76  	return cb()
    77  }
    78  
    79  // clone is used by the With* member functions to duplicate the current logger.
    80  func (ml *MockLogger) clone() *MockLogger {
    81  	result := &MockLogger{}
    82  	*result = *ml
    83  	return result
    84  }
    85  
    86  // WithJob duplicates the MockLogger and returns a new one tagged with the given
    87  // job ID.
    88  func (ml *MockLogger) WithJob(jobID string) TaggedLogger {
    89  	result := ml.clone()
    90  	result.Job = jobID
    91  	return result
    92  }
    93  
    94  // WithData duplicates the MockLogger and returns a new one tagged with the
    95  // given input data.
    96  func (ml *MockLogger) WithData(data []*common.Input) TaggedLogger {
    97  	result := ml.clone()
    98  	result.Data = data
    99  	return result
   100  }
   101  
   102  // WithUserCode duplicates the MockLogger and returns a new one tagged to
   103  // indicate that the log statements came from user code.
   104  func (ml *MockLogger) WithUserCode() TaggedLogger {
   105  	result := ml.clone()
   106  	result.UserCode = true
   107  	return result
   108  }
   109  
   110  // JobID returns the currently tagged job ID for the logger.  This is redundant
   111  // for MockLogger, as you can access ml.Job directly, but it is needed for the
   112  // TaggedLogger interface.
   113  func (ml *MockLogger) JobID() string {
   114  	return ml.Job
   115  }
   116  
   117  // Close is meant to be called to flush logs to object storage and return the
   118  // generated object, but this behavior is not implemented in MockLogger.
   119  func (ml *MockLogger) Close() (*pfs.Object, int64, error) {
   120  	// If you need an actual pfs.Object here, inherit and shadow this function
   121  	return nil, 0, nil
   122  }