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

     1  package datum
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/pachyderm/pachyderm/src/client"
     7  	"github.com/pachyderm/pachyderm/src/client/pfs"
     8  	"github.com/pachyderm/pachyderm/src/server/worker/common"
     9  )
    10  
    11  // A MockIterator can be used in tests to fulfill the Iterator interface without
    12  // having PFS running.  It returns Inputs based on the number range selected.
    13  type MockIterator struct {
    14  	start  int
    15  	end    int
    16  	index  int
    17  	inputs int
    18  }
    19  
    20  // MockIteratorOptions specifies the behavior of the MockIterator.
    21  type MockIteratorOptions struct {
    22  	Start  uint
    23  	Length uint
    24  	Inputs uint
    25  }
    26  
    27  // NewMockIterator constructs a mock iterator that will consistently return
    28  // Inputs based on the current index.
    29  func NewMockIterator(options *MockIteratorOptions) *MockIterator {
    30  	result := &MockIterator{
    31  		start:  int(options.Start),
    32  		end:    int(options.Start + options.Length),
    33  		index:  int(options.Start) - 1,
    34  		inputs: int(options.Inputs),
    35  	}
    36  
    37  	if result.inputs == 0 {
    38  		result.inputs = 1
    39  	}
    40  
    41  	return result
    42  }
    43  
    44  // Reset returns the current index to the start of the Iterator.
    45  func (mi *MockIterator) Reset() {
    46  	mi.index = mi.start - 1
    47  }
    48  
    49  // Len returns the number of items in the Iterator.
    50  func (mi *MockIterator) Len() int {
    51  	return int(mi.end - mi.start)
    52  }
    53  
    54  // Next steps the Iterator to the next item.
    55  func (mi *MockIterator) Next() bool {
    56  	if mi.index < mi.end {
    57  		mi.index++
    58  	}
    59  	return mi.index < mi.end
    60  }
    61  
    62  // Datum returns the set of Inputs for the current index of the Iterator.
    63  func (mi *MockIterator) Datum() []*common.Input {
    64  	return mi.DatumN(mi.index - mi.start)
    65  }
    66  
    67  // DatumN returns the set of Inputs for the selected index in the Iterator.
    68  func (mi *MockIterator) DatumN(index int) []*common.Input {
    69  	index = mi.start + index
    70  	if index < mi.start || index > mi.end {
    71  		panic(fmt.Sprintf("out of bounds access in MockIterator: index %d of range %d-%d", index, mi.start, mi.end))
    72  	}
    73  
    74  	result := []*common.Input{}
    75  	for i := 0; i < mi.inputs; i++ {
    76  		// Warning: this might break some assumptions about the format of this data
    77  		result = append(result, &common.Input{
    78  			FileInfo: &pfs.FileInfo{
    79  				File:      client.NewFile("dummy-repo", "dummy-commit", "/path"),
    80  				FileType:  pfs.FileType_FILE,
    81  				SizeBytes: uint64(index),
    82  				Hash:      []byte(fmt.Sprintf("%d", index)),
    83  			},
    84  			Name: fmt.Sprintf("source-%d", i),
    85  		})
    86  	}
    87  	return result
    88  }