github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+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  }