github.com/zaolin/u-root@v0.0.0-20200428085104-64aaafd46c6d/pkg/curl/mock_schemes.go (about)

     1  // Copyright 2017-2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package curl
     6  
     7  import (
     8  	"errors"
     9  	"io"
    10  	"net/url"
    11  	"path"
    12  	"strings"
    13  )
    14  
    15  // MockScheme is a Scheme mock for testing.
    16  type MockScheme struct {
    17  	// scheme is the scheme name.
    18  	Scheme string
    19  
    20  	// hosts is a map of host -> relative filename to host -> file contents.
    21  	hosts map[string]map[string]string
    22  
    23  	// numCalled is a map of URL string -> number of times Fetch has been
    24  	// called on that URL.
    25  	numCalled map[string]uint
    26  
    27  	// nextErr is the error to return for the next nextErrCount calls to
    28  	// Fetch. Note this introduces state into the MockScheme which is only
    29  	// okay in this scenario because MockScheme is only used for testing.
    30  	nextErr      error
    31  	nextErrCount int
    32  }
    33  
    34  // NewMockScheme creates a new MockScheme with the given scheme name.
    35  func NewMockScheme(scheme string) *MockScheme {
    36  	return &MockScheme{
    37  		Scheme:    scheme,
    38  		hosts:     make(map[string]map[string]string),
    39  		numCalled: make(map[string]uint),
    40  	}
    41  }
    42  
    43  // Add adds a file to the MockScheme
    44  func (m *MockScheme) Add(host string, p string, content string) {
    45  	_, ok := m.hosts[host]
    46  	if !ok {
    47  		m.hosts[host] = make(map[string]string)
    48  	}
    49  
    50  	m.hosts[host][path.Clean(p)] = content
    51  }
    52  
    53  // SetErr sets the error which is returned on the next count calls to Fetch.
    54  func (m *MockScheme) SetErr(err error, count int) {
    55  	m.nextErr = err
    56  	m.nextErrCount = count
    57  }
    58  
    59  // NumCalled returns how many times a url has been looked up.
    60  func (m *MockScheme) NumCalled(u *url.URL) uint {
    61  	url := u.String()
    62  	if c, ok := m.numCalled[url]; ok {
    63  		return c
    64  	}
    65  	return 0
    66  }
    67  
    68  var (
    69  	// ErrWrongScheme means the wrong mocked scheme was used.
    70  	ErrWrongScheme = errors.New("wrong scheme")
    71  	// ErrNoSuchHost means there is no host record in the mock.
    72  	ErrNoSuchHost = errors.New("no such host exists")
    73  	// ErrNoSuchFile means there is no file record in the mock.
    74  	ErrNoSuchFile = errors.New("no such file exists on this host")
    75  )
    76  
    77  // Fetch implements FileScheme.Fetch.
    78  func (m *MockScheme) Fetch(u *url.URL) (io.ReaderAt, error) {
    79  	url := u.String()
    80  	m.numCalled[url]++
    81  
    82  	if u.Scheme != m.Scheme {
    83  		return nil, ErrWrongScheme
    84  	}
    85  
    86  	if m.nextErrCount > 0 {
    87  		m.nextErrCount--
    88  		return nil, m.nextErr
    89  	}
    90  
    91  	files, ok := m.hosts[u.Host]
    92  	if !ok {
    93  		return nil, ErrNoSuchHost
    94  	}
    95  
    96  	content, ok := files[path.Clean(u.Path)]
    97  	if !ok {
    98  		return nil, ErrNoSuchFile
    99  	}
   100  	return strings.NewReader(content), nil
   101  }
   102  
   103  // MockSchemeRetryFilter is a Scheme mock for testing and has a method to
   104  // implement FileSchemeRetryFilter.
   105  type MockSchemeRetryFilter struct {
   106  	*MockScheme
   107  
   108  	retryFilter func(*url.URL, error) bool
   109  }
   110  
   111  // NewMockSchemeRetryFilter creates a new MockSchemeRetryFilter with the given
   112  // scheme name.
   113  func NewMockSchemeRetryFilter(scheme string) *MockSchemeRetryFilter {
   114  	return &MockSchemeRetryFilter{NewMockScheme(scheme), nil}
   115  }
   116  
   117  // RetryFilter implements FileSchemeRetryFilter.
   118  func (m *MockSchemeRetryFilter) RetryFilter(u *url.URL, err error) bool {
   119  	if m.retryFilter == nil {
   120  		return true
   121  	}
   122  	return m.retryFilter(u, err)
   123  }
   124  
   125  // SetRetryFilter sets the function to be used by the RetryFilter method.
   126  func (m *MockSchemeRetryFilter) SetRetryFilter(f func(*url.URL, error) bool) {
   127  	m.retryFilter = f
   128  }