gitee.com/mirrors_u-root/u-root@v7.0.0+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 }