kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package mockutil
15

16
import (
17
	"bytes"
18
	"io"
19
	"io/ioutil"
20
	"regexp"
21
)
22

23
// RegexMatcher is a gomock Matcher which matches strings against some
24
// given regex.
25
type RegexMatcher struct {
26
	expected *regexp.Regexp
27
}
28

29
// MatchRegex returns a new RegexMatcher which matches the expected regex.
30
func MatchRegex(expected string) *RegexMatcher {
31
	return &RegexMatcher{regexp.MustCompile(expected)}
32
}
33

34
// Matches returns true if x is a string which matches the expected regex.
35
func (m *RegexMatcher) Matches(x interface{}) bool {
36
	s, ok := x.(string)
37
	if !ok {
38
		return false
39
	}
40
	return m.expected.MatchString(s)
41
}
42

43
func (m *RegexMatcher) String() string {
44
	return m.expected.String()
45
}
46

47
// ReaderMatcher is a gomock Matcher which matches io.Readers which produce some
48
// given bytes.
49
type ReaderMatcher struct {
50
	expected []byte
51
}
52

53
// MatchReader returns a new ReaderMatcher which matches expected.
54
func MatchReader(expected []byte) *ReaderMatcher {
55
	return &ReaderMatcher{expected}
56
}
57

58
// Matches returns true if x is an io.Reader which contains the expected bytes.
59
func (m *ReaderMatcher) Matches(x interface{}) bool {
60
	r, ok := x.(io.Reader)
61
	if !ok {
62
		return false
63
	}
64
	b, err := ioutil.ReadAll(r)
65
	if err != nil {
66
		panic(err)
67
	}
68
	return bytes.Compare(m.expected, b) == 0
69
}
70

71
func (m *ReaderMatcher) String() string {
72
	return string(m.expected)
73
}
74

75
// WriterMatcher is a gomock Matcher which matches any io.Writer, with the
76
// side-effect of writing some given bytes.
77
type WriterMatcher struct {
78
	b []byte
79
}
80

81
// MatchWriter returns a new WriterMatcher which write b to any io.Writer passed
82
// to Matches.
83
func MatchWriter(b []byte) *WriterMatcher {
84
	return &WriterMatcher{b}
85
}
86

87
// Matches writes given bytes to x.
88
func (m *WriterMatcher) Matches(x interface{}) bool {
89
	w, ok := x.(io.Writer)
90
	if !ok {
91
		return false
92
	}
93
	if _, err := w.Write(m.b); err != nil {
94
		panic(err)
95
	}
96
	return true
97
}
98

99
func (m *WriterMatcher) String() string {
100
	return "WriterMatcher"
101
}
102

103
// WriterAtMatcher is a gomock Matcher which matches any io.WriterAt, with the
104
// side-effect of writing some give bytes.
105
type WriterAtMatcher struct {
106
	b []byte
107
}
108

109
// MatchWriterAt returns a new WriterAtMatcher which writes b to any io.WriterAt passed
110
// to Matches.
111
func MatchWriterAt(b []byte) *WriterAtMatcher {
112
	return &WriterAtMatcher{b}
113
}
114

115
// Matches writes given bytes to x.
116
func (m *WriterAtMatcher) Matches(x interface{}) bool {
117
	w, ok := x.(io.WriterAt)
118
	if !ok {
119
		return false
120
	}
121
	if _, err := w.WriteAt(m.b, 0); err != nil {
122
		panic(err)
123
	}
124
	return true
125
}
126

127
func (m *WriterAtMatcher) String() string {
128
	return "WriterAtMatcher"
129
}
130