github.com/pachyderm/pachyderm@v1.13.4/src/server/pkg/testpachd/mock_env.go (about)

     1  package testpachd
     2  
     3  import (
     4  	"context"
     5  
     6  	"golang.org/x/sync/errgroup"
     7  
     8  	"github.com/pachyderm/pachyderm/src/client"
     9  	"github.com/pachyderm/pachyderm/src/server/pkg/testetcd"
    10  )
    11  
    12  // MockEnv contains the basic setup for running end-to-end pachyderm tests
    13  // entirely locally within the test process. It provides a temporary directory
    14  // for storing data, an embedded etcd server with a connected client, as well as
    15  // a local mock pachd instance which allows a test to hook into any pachd calls.
    16  type MockEnv struct {
    17  	testetcd.Env
    18  	MockPachd  *MockPachd
    19  	PachClient *client.APIClient
    20  }
    21  
    22  // WithMockEnv sets up a MockEnv structure, passes it to the provided callback,
    23  // then cleans up everything in the environment, regardless of if an assertion
    24  // fails.
    25  func WithMockEnv(cb func(*MockEnv) error) error {
    26  	return testetcd.WithEnv(func(etcdEnv *testetcd.Env) (err error) {
    27  		// Use an error group with a cancelable context to supervise every component
    28  		// and cancel everything if one fails
    29  		ctx, cancel := context.WithCancel(etcdEnv.Context)
    30  		defer cancel()
    31  		eg, ctx := errgroup.WithContext(ctx)
    32  
    33  		mockEnv := &MockEnv{Env: *etcdEnv}
    34  		mockEnv.Context = ctx
    35  
    36  		// Cleanup any state when we return
    37  		defer func() {
    38  			saveErr := func(e error) error {
    39  				if e != nil && err == nil {
    40  					err = e
    41  				}
    42  				return e
    43  			}
    44  
    45  			if mockEnv.PachClient != nil {
    46  				saveErr(mockEnv.PachClient.Close())
    47  			}
    48  
    49  			if mockEnv.MockPachd != nil {
    50  				saveErr(mockEnv.MockPachd.Close())
    51  			}
    52  
    53  			cancel()
    54  			saveErr(eg.Wait())
    55  		}()
    56  
    57  		mockEnv.MockPachd, err = NewMockPachd(mockEnv.Context)
    58  		if err != nil {
    59  			return err
    60  		}
    61  
    62  		eg.Go(func() error {
    63  			return errorWait(ctx, mockEnv.MockPachd.Err())
    64  		})
    65  		mockEnv.PachClient, err = client.NewFromAddress(mockEnv.MockPachd.Addr.String())
    66  		if err != nil {
    67  			return err
    68  		}
    69  
    70  		// TODO: supervise the PachClient connection and error the errgroup if they
    71  		// go down
    72  
    73  		return cb(mockEnv)
    74  	})
    75  }
    76  
    77  func errorWait(ctx context.Context, errChan <-chan error) error {
    78  	select {
    79  	case <-ctx.Done():
    80  		return nil
    81  	case err := <-errChan:
    82  		return err
    83  	}
    84  }