cubefs

1
/*
2
Package mocks provides mocks that can be used for testing applications
3
that use Sarama. The mock types provided by this package implement the
4
interfaces Sarama exports, so you can use them for dependency injection
5
in your tests.
6

7
All mock instances require you to set expectations on them before you
8
can use them. It will determine how the mock will behave. If an
9
expectation is not met, it will make your test fail.
10

11
NOTE: this package currently does not fall under the API stability
12
guarantee of Sarama as it is still considered experimental.
13
*/
14
package mocks
15

16
import (
17
	"errors"
18
	"fmt"
19

20
	"github.com/Shopify/sarama"
21
)
22

23
// ErrorReporter is a simple interface that includes the testing.T methods we use to report
24
// expectation violations when using the mock objects.
25
type ErrorReporter interface {
26
	Errorf(string, ...interface{})
27
}
28

29
// ValueChecker is a function type to be set in each expectation of the producer mocks
30
// to check the value passed.
31
type ValueChecker func(val []byte) error
32

33
// MessageChecker is a function type to be set in each expectation of the producer mocks
34
// to check the message passed.
35
type MessageChecker func(*sarama.ProducerMessage) error
36

37
// messageValueChecker wraps a ValueChecker into a MessageChecker.
38
// Failure to encode the message value will return an error and not call
39
// the wrapped ValueChecker.
40
func messageValueChecker(f ValueChecker) MessageChecker {
41
	if f == nil {
42
		return nil
43
	}
44
	return func(msg *sarama.ProducerMessage) error {
45
		val, err := msg.Value.Encode()
46
		if err != nil {
47
			return fmt.Errorf("Input message encoding failed: %w", err)
48
		}
49
		return f(val)
50
	}
51
}
52

53
var (
54
	errProduceSuccess              error = nil
55
	errOutOfExpectations                 = errors.New("No more expectations set on mock")
56
	errPartitionConsumerNotStarted       = errors.New("The partition consumer was never started")
57
)
58

59
const AnyOffset int64 = -1000
60

61
type producerExpectation struct {
62
	Result        error
63
	CheckFunction MessageChecker
64
}
65

66
// TopicConfig describes a mock topic structure for the mock producers’ partitioning needs.
67
type TopicConfig struct {
68
	overridePartitions map[string]int32
69
	defaultPartitions  int32
70
}
71

72
// NewTopicConfig makes a configuration which defaults to 32 partitions for every topic.
73
func NewTopicConfig() *TopicConfig {
74
	return &TopicConfig{
75
		overridePartitions: make(map[string]int32, 0),
76
		defaultPartitions:  32,
77
	}
78
}
79

80
// SetDefaultPartitions sets the number of partitions any topic not explicitly configured otherwise
81
// (by SetPartitions) will have from the perspective of created partitioners.
82
func (pc *TopicConfig) SetDefaultPartitions(n int32) {
83
	pc.defaultPartitions = n
84
}
85

86
// SetPartitions sets the number of partitions the partitioners will see for specific topics. This
87
// only applies to messages produced after setting them.
88
func (pc *TopicConfig) SetPartitions(partitions map[string]int32) {
89
	for p, n := range partitions {
90
		pc.overridePartitions[p] = n
91
	}
92
}
93

94
func (pc *TopicConfig) partitions(topic string) int32 {
95
	if n, found := pc.overridePartitions[topic]; found {
96
		return n
97
	}
98
	return pc.defaultPartitions
99
}
100

101
// NewTestConfig returns a config meant to be used by tests.
102
// Due to inconsistencies with the request versions the clients send using the default Kafka version
103
// and the response versions our mocks use, we default to the minimum Kafka version in most tests
104
func NewTestConfig() *sarama.Config {
105
	config := sarama.NewConfig()
106
	config.Version = sarama.MinVersion
107
	return config
108
}
109