Extend client mock with ability to test publish, and a few useful method like SetResponse and SetSubscriber (#2375)

Co-authored-by: Hunyadvári Péter <peter.hunyadvari@vcc.live>

plugins/client/mock/context.go

@@ -6,11 +6,22 @@ import (
 
 type responseKey struct{}
 
-func fromContext(ctx context.Context) (map[string][]MockResponse, bool) {
+func responseFromContext(ctx context.Context) (map[string][]MockResponse, bool) {
 	r, ok := ctx.Value(responseKey{}).(map[string][]MockResponse)
 	return r, ok
 }
 
-func newContext(ctx context.Context, r map[string][]MockResponse) context.Context {
+func newResponseContext(ctx context.Context, r map[string][]MockResponse) context.Context {
 	return context.WithValue(ctx, responseKey{}, r)
 }
+
+type subscriberKey struct{}
+
+func subscriberFromContext(ctx context.Context) (map[string]MockSubscriber, bool) {
+	r, ok := ctx.Value(subscriberKey{}).(map[string]MockSubscriber)
+	return r, ok
+}
+
+func newSubscriberContext(ctx context.Context, r map[string]MockSubscriber) context.Context {
+	return context.WithValue(ctx, subscriberKey{}, r)
+}

plugins/client/mock/mock.go

@@ -21,12 +21,15 @@ type MockResponse struct {
 	Error    error
 }
 
+type MockSubscriber func(client.Message) error
+
 type MockClient struct {
 	Client client.Client
 	Opts   client.Options
 
 	sync.Mutex
-	Response map[string][]MockResponse
+	Response   map[string][]MockResponse
+	Subscriber map[string]MockSubscriber
 }
 
 func (m *MockClient) Init(opts ...client.Option) error {
@@ -37,7 +40,7 @@ func (m *MockClient) Init(opts ...client.Option) error {
 		opt(&m.Opts)
 	}
 
-	r, ok := fromContext(m.Opts.Context)
+	r, ok := responseFromContext(m.Opts.Context)
 	if !ok {
 		r = make(map[string][]MockResponse)
 	}
@@ -121,6 +124,12 @@ func (m *MockClient) Stream(ctx context.Context, req client.Request, opts ...cli
 }
 
 func (m *MockClient) Publish(ctx context.Context, p client.Message, opts ...client.PublishOption) error {
+	m.Lock()
+	defer m.Unlock()
+
+	if s, ok := m.Subscriber[p.Topic()]; ok {
+		return s(p)
+	}
 	return nil
 }
 
@@ -128,6 +137,34 @@ func (m *MockClient) String() string {
 	return "mock"
 }
 
+func (m *MockClient) SetResponse(service string, response MockResponse) {
+	m.Lock()
+	defer m.Unlock()
+
+	endpoints, ok := m.Response[service]
+	if !ok {
+		// new service
+		m.Response[service] = []MockResponse{response}
+		return
+	}
+	for i, r := range endpoints {
+		if r.Endpoint == response.Endpoint {
+			// update service endpoint
+			m.Response[service][i] = response
+			return
+		}
+	}
+	// add new endpoint for service
+	m.Response[service] = append(m.Response[service], response)
+}
+
+func (m *MockClient) SetSubscriber(topic string, subscriber MockSubscriber) {
+	m.Lock()
+	defer m.Unlock()
+
+	m.Subscriber[topic] = subscriber
+}
+
 func NewClient(opts ...client.Option) *MockClient {
 	options := client.Options{
 		Context: context.TODO(),
@@ -137,14 +174,20 @@ func NewClient(opts ...client.Option) *MockClient {
 		opt(&options)
 	}
 
-	r, ok := fromContext(options.Context)
+	r, ok := responseFromContext(options.Context)
 	if !ok {
 		r = make(map[string][]MockResponse)
 	}
 
+	s, ok := subscriberFromContext(options.Context)
+	if !ok {
+		s = make(map[string]MockSubscriber)
+	}
+
 	return &MockClient{
-		Client:   client.DefaultClient,
-		Opts:     options,
-		Response: r,
+		Client:     client.DefaultClient,
+		Opts:       options,
+		Response:   r,
+		Subscriber: s,
 	}
 }

plugins/client/mock/options.go

@@ -7,11 +7,23 @@ import (
 // Response sets the response methods for a service
 func Response(service string, response []MockResponse) client.Option {
 	return func(o *client.Options) {
-		r, ok := fromContext(o.Context)
+		r, ok := responseFromContext(o.Context)
 		if !ok {
 			r = make(map[string][]MockResponse)
 		}
 		r[service] = response
-		o.Context = newContext(o.Context, r)
+		o.Context = newResponseContext(o.Context, r)
+	}
+}
+
+// Subscriber sets the subscribers service
+func Subscriber(topic string, subscriber MockSubscriber) client.Option {
+	return func(o *client.Options) {
+		r, ok := subscriberFromContext(o.Context)
+		if !ok {
+			r = make(map[string]MockSubscriber)
+		}
+		r[topic] = subscriber
+		o.Context = newSubscriberContext(o.Context, r)
 	}
 }