github.com/manicqin/nomad@v0.9.5/plugins/device/mock.go (about)

     1  package device
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/hashicorp/nomad/plugins/base"
     8  )
     9  
    10  type FingerprintFn func(context.Context) (<-chan *FingerprintResponse, error)
    11  type ReserveFn func([]string) (*ContainerReservation, error)
    12  type StatsFn func(context.Context, time.Duration) (<-chan *StatsResponse, error)
    13  
    14  // MockDevicePlugin is used for testing.
    15  // Each function can be set as a closure to make assertions about how data
    16  // is passed through the base plugin layer.
    17  type MockDevicePlugin struct {
    18  	*base.MockPlugin
    19  	FingerprintF FingerprintFn
    20  	ReserveF     ReserveFn
    21  	StatsF       StatsFn
    22  }
    23  
    24  func (p *MockDevicePlugin) Fingerprint(ctx context.Context) (<-chan *FingerprintResponse, error) {
    25  	return p.FingerprintF(ctx)
    26  }
    27  
    28  func (p *MockDevicePlugin) Reserve(devices []string) (*ContainerReservation, error) {
    29  	return p.ReserveF(devices)
    30  }
    31  
    32  func (p *MockDevicePlugin) Stats(ctx context.Context, interval time.Duration) (<-chan *StatsResponse, error) {
    33  	return p.StatsF(ctx, interval)
    34  }
    35  
    36  // Below are static implementations of the device functions
    37  
    38  // StaticFingerprinter fingerprints the passed devices just once
    39  func StaticFingerprinter(devices []*DeviceGroup) FingerprintFn {
    40  	return func(_ context.Context) (<-chan *FingerprintResponse, error) {
    41  		outCh := make(chan *FingerprintResponse, 1)
    42  		outCh <- &FingerprintResponse{
    43  			Devices: devices,
    44  		}
    45  		return outCh, nil
    46  	}
    47  }
    48  
    49  // ErrorChFingerprinter returns an error fingerprinting over the channel
    50  func ErrorChFingerprinter(err error) FingerprintFn {
    51  	return func(_ context.Context) (<-chan *FingerprintResponse, error) {
    52  		outCh := make(chan *FingerprintResponse, 1)
    53  		outCh <- &FingerprintResponse{
    54  			Error: err,
    55  		}
    56  		return outCh, nil
    57  	}
    58  }
    59  
    60  // StaticReserve returns the passed container reservation
    61  func StaticReserve(out *ContainerReservation) ReserveFn {
    62  	return func(_ []string) (*ContainerReservation, error) {
    63  		return out, nil
    64  	}
    65  }
    66  
    67  // ErrorReserve returns the passed error
    68  func ErrorReserve(err error) ReserveFn {
    69  	return func(_ []string) (*ContainerReservation, error) {
    70  		return nil, err
    71  	}
    72  }
    73  
    74  // StaticStats returns the passed statistics
    75  func StaticStats(out []*DeviceGroupStats) StatsFn {
    76  	return func(ctx context.Context, intv time.Duration) (<-chan *StatsResponse, error) {
    77  		outCh := make(chan *StatsResponse, 1)
    78  
    79  		go func() {
    80  			ticker := time.NewTimer(0)
    81  			for {
    82  				select {
    83  				case <-ctx.Done():
    84  					return
    85  				case <-ticker.C:
    86  					ticker.Reset(intv)
    87  				}
    88  
    89  				outCh <- &StatsResponse{
    90  					Groups: out,
    91  				}
    92  			}
    93  		}()
    94  
    95  		return outCh, nil
    96  	}
    97  }
    98  
    99  // ErrorChStats returns an error collecting stats over the channel
   100  func ErrorChStats(err error) StatsFn {
   101  	return func(_ context.Context, _ time.Duration) (<-chan *StatsResponse, error) {
   102  		outCh := make(chan *StatsResponse, 1)
   103  		outCh <- &StatsResponse{
   104  			Error: err,
   105  		}
   106  		return outCh, nil
   107  	}
   108  }