github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/dm/mock_message_agent.go (about)

     1  // Copyright 2022 PingCAP, 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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package dm
    15  
    16  import (
    17  	"context"
    18  	"encoding/json"
    19  	"sync"
    20  
    21  	"github.com/stretchr/testify/mock"
    22  )
    23  
    24  // MockMessageAgent implement MessageAgent
    25  type MockMessageAgent struct {
    26  	sync.Mutex
    27  	mock.Mock
    28  }
    29  
    30  var _ MessageAgent = &MockMessageAgent{}
    31  
    32  // GenerateTopic generate mock message topic.
    33  func GenerateTopic(senderID, receiverID string) string {
    34  	return generateTopic(senderID, receiverID)
    35  }
    36  
    37  // GenerateResponse generate mock response message.
    38  func GenerateResponse(id messageID, command string, msg interface{}) interface{} {
    39  	resp := message{ID: id, Type: responseTp, Command: command, Payload: msg}
    40  	// nolint:errcheck
    41  	bytes, _ := json.Marshal(resp)
    42  	var resp2 message
    43  	// nolint:errcheck
    44  	json.Unmarshal(bytes, &resp2)
    45  	return &resp2
    46  }
    47  
    48  // Tick implement MessageAgent.Tick.
    49  func (m *MockMessageAgent) Tick(ctx context.Context) error { return nil }
    50  
    51  // Close implement MessageAgent.Close.
    52  func (m *MockMessageAgent) Close(ctx context.Context) error { return nil }
    53  
    54  // UpdateClient implement MessageAgent.UpdateClient.
    55  func (m *MockMessageAgent) UpdateClient(clientID string, client client) error { return nil }
    56  
    57  // RemoveClient implement MessageAgent.RemoveClient.
    58  func (m *MockMessageAgent) RemoveClient(clientID string) error { return nil }
    59  
    60  // SendMessage implement MessageAgent.SendMessage.
    61  func (m *MockMessageAgent) SendMessage(ctx context.Context, clientID string, command string, msg interface{}) error {
    62  	m.Lock()
    63  	defer m.Unlock()
    64  	return m.Called().Error(0)
    65  }
    66  
    67  // SendRequest implement MessageAgent.SendRequest.
    68  func (m *MockMessageAgent) SendRequest(ctx context.Context, clientID string, command string, req interface{}) (interface{}, error) {
    69  	m.Lock()
    70  	defer m.Unlock()
    71  	args := m.Called(ctx, clientID, command, req)
    72  	return args.Get(0), args.Error(1)
    73  }