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