github.com/alkemics/goflow@v0.2.1/wrappers/mockingjay/mockingjay.go (about)

     1  package mockingjay
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"github.com/alkemics/goflow"
     9  	"github.com/alkemics/goflow/wrappers/ctx"
    10  )
    11  
    12  type ContextKeyType string
    13  
    14  const ContextKey ContextKeyType = "mockingjay"
    15  
    16  func WithMock(ctx context.Context, nodeID string, values ...interface{}) context.Context {
    17  	m, ok := ctx.Value(ContextKey).(map[string][]interface{})
    18  	if !ok || m == nil {
    19  		m = make(map[string][]interface{})
    20  	}
    21  
    22  	m[nodeID] = values
    23  	return context.WithValue(ctx, ContextKey, m)
    24  }
    25  
    26  // Mock is a goflow.NodeWrapper that lets you mock nodes at run time.
    27  //
    28  // The mock will extract the value from the context at ContextKey as a
    29  // map[string][]interface{}, looking in the map via the node id to find
    30  // the mocked returns. If the node is not mocked, it is executed normally.
    31  //
    32  // The WithMock helper function is provided to make it easier to mock  node:
    33  //
    34  //     ctx := mockingjay.WithMock(ctx, "myNode", 42)
    35  //     graph.Run(ctx, ...)
    36  func Mock(_ func(interface{}) error, node goflow.NodeRenderer) (goflow.NodeRenderer, error) {
    37  	return mocker{
    38  		NodeRenderer: ctx.Injector{
    39  			NodeRenderer: node,
    40  		},
    41  	}, nil
    42  }
    43  
    44  type mocker struct {
    45  	goflow.NodeRenderer
    46  }
    47  
    48  func (m mocker) Run(inputs, outputs []goflow.Field) (string, error) {
    49  	sub, err := m.NodeRenderer.Run(inputs, outputs)
    50  	if err != nil {
    51  		return "", err
    52  	}
    53  
    54  	mocked := make([]string, len(outputs))
    55  	for i, output := range outputs {
    56  		mocked[i] = fmt.Sprintf("%s = _mock[%d].(%s)", output.Name, i, output.Type)
    57  	}
    58  
    59  	return fmt.Sprintf(`
    60  var _mock []interface{}
    61  if _mocks, ok := ctx.Value(mockingjay.ContextKey).(map[string][]interface{}); ok && _mocks != nil {
    62  	m, ok := _mocks["%s"]
    63  	if ok {
    64  		_mock = m
    65  	}
    66  }
    67  
    68  if _mock != nil {
    69  	%s
    70  } else {
    71  	%s
    72  }
    73  `,
    74  		m.ID(),
    75  		strings.Join(mocked, "\n"),
    76  		sub,
    77  	), nil
    78  }