github.com/system-transparency/u-root@v6.0.1-0.20190919065413-ed07a650de4c+incompatible/pkg/urlfetch/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 urlfetch
     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  
    28  // NewMockScheme creates a new MockScheme with the given scheme name.
    29  func NewMockScheme(scheme string) *MockScheme {
    30  	return &MockScheme{
    31  		Scheme:    scheme,
    32  		hosts:     make(map[string]map[string]string),
    33  		numCalled: make(map[string]uint),
    34  	}
    35  }
    36  
    37  // Add adds a file to the MockScheme
    38  func (m *MockScheme) Add(host string, p string, content string) {
    39  	_, ok := m.hosts[host]
    40  	if !ok {
    41  		m.hosts[host] = make(map[string]string)
    42  	}
    43  
    44  	m.hosts[host][path.Clean(p)] = content
    45  }
    46  
    47  // NumCalled returns how many times a url has been looked up.
    48  func (m *MockScheme) NumCalled(u *url.URL) uint {
    49  	url := u.String()
    50  	if c, ok := m.numCalled[url]; ok {
    51  		return c
    52  	}
    53  	return 0
    54  }
    55  
    56  var (
    57  	// ErrWrongScheme means the wrong mocked scheme was used.
    58  	ErrWrongScheme = errors.New("wrong scheme")
    59  	// ErrNoSuchHost means there is no host record in the mock.
    60  	ErrNoSuchHost = errors.New("no such host exists")
    61  	// ErrNoSuchFile means there is no file record in the mock.
    62  	ErrNoSuchFile = errors.New("no such file exists on this host")
    63  )
    64  
    65  // Fetch implements FileScheme.Fetch.
    66  func (m *MockScheme) Fetch(u *url.URL) (io.ReaderAt, error) {
    67  	url := u.String()
    68  	if _, ok := m.numCalled[url]; ok {
    69  		m.numCalled[url]++
    70  	} else {
    71  		m.numCalled[url] = 1
    72  	}
    73  
    74  	if u.Scheme != m.Scheme {
    75  		return nil, ErrWrongScheme
    76  	}
    77  
    78  	files, ok := m.hosts[u.Host]
    79  	if !ok {
    80  		return nil, ErrNoSuchHost
    81  	}
    82  
    83  	content, ok := files[path.Clean(u.Path)]
    84  	if !ok {
    85  		return nil, ErrNoSuchFile
    86  	}
    87  	return strings.NewReader(content), nil
    88  }