github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/dm/master/workerrpc/workerrpc_test.go (about)

     1  // Copyright 2020 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 workerrpc
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  	"time"
    20  
    21  	"github.com/golang/mock/gomock"
    22  	. "github.com/pingcap/check"
    23  	"github.com/pingcap/errors"
    24  	"github.com/pingcap/tiflow/dm/config/security"
    25  	"github.com/pingcap/tiflow/dm/pb"
    26  	"github.com/pingcap/tiflow/dm/pbmock"
    27  	"github.com/pingcap/tiflow/dm/pkg/terror"
    28  	"github.com/tikv/pd/pkg/utils/tempurl"
    29  )
    30  
    31  var _ = Suite(&testWorkerRPCSuite{})
    32  
    33  type testWorkerRPCSuite struct{}
    34  
    35  func TestWorkerRPC(t *testing.T) {
    36  	TestingT(t)
    37  }
    38  
    39  func (t *testWorkerRPCSuite) TestGRPCClient(c *C) {
    40  	timeout := 3 * time.Second
    41  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
    42  	defer cancel()
    43  
    44  	ctrl := gomock.NewController(c)
    45  	defer ctrl.Finish()
    46  
    47  	// get a random address for DM-worker
    48  	addr := tempurl.Alloc()[len("http://"):]
    49  	// NOTE: we don't wait for the gRPC connection establish now, in other words no need to wait for the DM-worker instance become online.
    50  	rpcCli, err := NewGRPCClient(addr, security.Security{})
    51  	c.Assert(err, IsNil)
    52  
    53  	// replace the underlying DM-worker client.
    54  	workerCli := pbmock.NewMockWorkerClient(ctrl)
    55  	rpcCli.client = workerCli
    56  
    57  	reqs := []*Request{
    58  		{
    59  			Type:        CmdQueryStatus,
    60  			QueryStatus: &pb.QueryStatusRequest{Name: "test"},
    61  		},
    62  		{
    63  			Type:       CmdPurgeRelay,
    64  			PurgeRelay: &pb.PurgeRelayRequest{Inactive: true},
    65  		},
    66  		{
    67  			Type:          CmdOperateSchema,
    68  			OperateSchema: &pb.OperateWorkerSchemaRequest{Op: pb.SchemaOp_SetSchema},
    69  		},
    70  		{
    71  			Type:          CmdOperateV1Meta,
    72  			OperateV1Meta: &pb.OperateV1MetaRequest{Op: pb.V1MetaOp_RemoveV1Meta},
    73  		},
    74  		{
    75  			Type:        CmdHandleError,
    76  			HandleError: &pb.HandleWorkerErrorRequest{Op: pb.ErrorOp_Replace},
    77  		},
    78  	}
    79  
    80  	workerCli.EXPECT().QueryStatus(gomock.Any(), reqs[0].QueryStatus)
    81  	workerCli.EXPECT().PurgeRelay(gomock.Any(), reqs[1].PurgeRelay)
    82  	workerCli.EXPECT().OperateSchema(gomock.Any(), reqs[2].OperateSchema)
    83  	workerCli.EXPECT().OperateV1Meta(gomock.Any(), reqs[3].OperateV1Meta)
    84  	workerCli.EXPECT().HandleError(gomock.Any(), reqs[4].HandleError)
    85  
    86  	// others cmds are not supported.
    87  	// NOTE: update the end cmd in the below `for` loop when adding new cmds.
    88  OUTER:
    89  	for cmd := CmdStartSubTask; cmd <= CmdHandleError; cmd++ {
    90  		for _, req := range reqs {
    91  			if req.Type == cmd {
    92  				// supported cmd
    93  				_, err = rpcCli.SendRequest(ctx, req, timeout)
    94  				c.Assert(err, IsNil)
    95  				continue OUTER
    96  			}
    97  		}
    98  		_, err = rpcCli.SendRequest(ctx, &Request{Type: cmd}, timeout)
    99  		c.Assert(terror.ErrMasterGRPCInvalidReqType.Equal(err), IsTrue)
   100  	}
   101  
   102  	// got an error from the underlying RPC.
   103  	err2 := errors.New("mock error")
   104  	workerCli.EXPECT().QueryStatus(gomock.Any(), reqs[0].QueryStatus).Return(nil, err2)
   105  	_, err = rpcCli.SendRequest(ctx, reqs[0], timeout)
   106  	c.Assert(terror.ErrMasterGRPCRequestError.Equal(err), IsTrue)
   107  	c.Assert(errors.Cause(err), Equals, err2)
   108  
   109  	// close the cli.
   110  	c.Assert(rpcCli.Close(), IsNil)
   111  
   112  	// can't send request any more.
   113  	_, err = rpcCli.SendRequest(ctx, reqs[0], timeout)
   114  	c.Assert(terror.ErrMasterGRPCSendOnCloseConn.Equal(err), IsTrue)
   115  }