github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/client/test_util.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 client
    15  
    16  import (
    17  	"bytes"
    18  	"encoding/json"
    19  	"fmt"
    20  	"reflect"
    21  
    22  	"github.com/golang/mock/gomock"
    23  	"github.com/pingcap/tiflow/engine/enginepb"
    24  )
    25  
    26  type preDispatchArgsMatcher struct {
    27  	args *DispatchTaskArgs
    28  }
    29  
    30  func matchPreDispatchArgs(args *DispatchTaskArgs) gomock.Matcher {
    31  	return &preDispatchArgsMatcher{args: args}
    32  }
    33  
    34  func (m *preDispatchArgsMatcher) Matches(x interface{}) bool {
    35  	// TODO match ProjectInfo
    36  	req, ok := x.(*enginepb.PreDispatchTaskRequest)
    37  	if !ok {
    38  		return false
    39  	}
    40  
    41  	if !(m.args.WorkerID == req.GetWorkerId() &&
    42  		m.args.MasterID == req.GetMasterId() &&
    43  		m.args.WorkerType == req.GetTaskTypeId()) {
    44  
    45  		return false
    46  	}
    47  
    48  	if !bytes.Equal(m.args.WorkerConfig, req.GetTaskConfig()) {
    49  		return false
    50  	}
    51  
    52  	return true
    53  }
    54  
    55  func (m *preDispatchArgsMatcher) String() string {
    56  	jsonBytes, err := json.Marshal(m.args)
    57  	if err != nil {
    58  		panic(err)
    59  	}
    60  
    61  	return fmt.Sprintf("PreDispatchRequest Matches Args %s", string(jsonBytes))
    62  }
    63  
    64  type confirmDispatchMatcher struct {
    65  	// requestID is a pointer to a string,
    66  	// because the pointed string will be changed
    67  	// after the matcher has been created.
    68  	requestID *string
    69  
    70  	workerID string
    71  }
    72  
    73  func matchConfirmDispatch(requestID *string, workerID string) gomock.Matcher {
    74  	return &confirmDispatchMatcher{
    75  		requestID: requestID,
    76  		workerID:  workerID,
    77  	}
    78  }
    79  
    80  func (m *confirmDispatchMatcher) Matches(x interface{}) bool {
    81  	req, ok := x.(*enginepb.ConfirmDispatchTaskRequest)
    82  	if !ok {
    83  		return false
    84  	}
    85  
    86  	return reflect.DeepEqual(req, &enginepb.ConfirmDispatchTaskRequest{
    87  		WorkerId:  m.workerID,
    88  		RequestId: *m.requestID,
    89  	})
    90  }
    91  
    92  func (m *confirmDispatchMatcher) String() string {
    93  	expected := &enginepb.ConfirmDispatchTaskRequest{
    94  		WorkerId:  m.workerID,
    95  		RequestId: *m.requestID,
    96  	}
    97  
    98  	return fmt.Sprintf("ConfirmDispatchRequest Matches Args %s", expected.String())
    99  }