github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/internal/master/worker_creator_test.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 master
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/golang/mock/gomock"
    21  	"github.com/pingcap/log"
    22  	pb "github.com/pingcap/tiflow/engine/enginepb"
    23  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    24  	"github.com/pingcap/tiflow/engine/model"
    25  	"github.com/pingcap/tiflow/engine/pkg/client"
    26  	resModel "github.com/pingcap/tiflow/engine/pkg/externalresource/model"
    27  	pkgOrm "github.com/pingcap/tiflow/engine/pkg/orm"
    28  	"github.com/pingcap/tiflow/engine/pkg/tenant"
    29  	"github.com/pingcap/tiflow/pkg/label"
    30  	"github.com/stretchr/testify/require"
    31  	"go.uber.org/atomic"
    32  )
    33  
    34  type mockEpochGen struct {
    35  	pkgOrm.Client
    36  
    37  	nextEpoch int64
    38  }
    39  
    40  func (g *mockEpochGen) GenEpoch(ctx context.Context) (int64, error) {
    41  	return g.nextEpoch, nil
    42  }
    43  
    44  type workerCreatorTestHelper struct {
    45  	ExecutorGroup      *client.MockExecutorGroup
    46  	ServerMasterClient *client.MockServerMasterClient
    47  	FrameMetaClient    *mockEpochGen
    48  
    49  	Creator *WorkerCreator
    50  }
    51  
    52  func newWorkerCreatorTestHelper(
    53  	t *testing.T,
    54  	jobID frameModel.MasterID,
    55  	hooks *WorkerCreationHooks,
    56  	inheritedSelectors []*label.Selector,
    57  ) *workerCreatorTestHelper {
    58  	ret := &workerCreatorTestHelper{
    59  		ExecutorGroup:      client.NewMockExecutorGroup(),
    60  		ServerMasterClient: client.NewMockServerMasterClient(gomock.NewController(t)),
    61  		FrameMetaClient:    &mockEpochGen{nextEpoch: 1},
    62  	}
    63  
    64  	ret.Creator = NewWorkerCreatorBuilder().
    65  		WithMasterID(jobID).
    66  		WithLogger(log.L()).
    67  		WithInheritedSelectors(inheritedSelectors...).
    68  		WithExecutorGroup(ret.ExecutorGroup).
    69  		WithFrameMetaClient(ret.FrameMetaClient).
    70  		WithServerMasterClient(ret.ServerMasterClient).
    71  		WithHooks(hooks).Build()
    72  	return ret
    73  }
    74  
    75  func TestCreateWorkerNormal(t *testing.T) {
    76  	t.Parallel()
    77  
    78  	inheritedSelectors := []*label.Selector{
    79  		{
    80  			Key:    "label1",
    81  			Target: "value1",
    82  			Op:     label.OpEq,
    83  		},
    84  		{
    85  			Key:    "label2",
    86  			Target: "value2",
    87  			Op:     label.OpNeq,
    88  		},
    89  	}
    90  	additionalSelectors := []*label.Selector{
    91  		{
    92  			Key:    "label3",
    93  			Target: "value3",
    94  			Op:     label.OpEq,
    95  		},
    96  		{
    97  			Key:    "label4",
    98  			Target: "value4",
    99  			Op:     label.OpNeq,
   100  		},
   101  	}
   102  	var callBackInvoked atomic.Bool
   103  	helper := newWorkerCreatorTestHelper(t, "job-1",
   104  		&WorkerCreationHooks{
   105  			BeforeStartingWorker: func(
   106  				workerID frameModel.WorkerID,
   107  				executorID model.ExecutorID,
   108  				epoch frameModel.Epoch,
   109  			) {
   110  				require.Equal(t, "worker-1", workerID)
   111  				require.Equal(t, "executor-1", executorID)
   112  				require.Equal(t, int64(1), epoch)
   113  				require.False(t, callBackInvoked.Swap(true))
   114  			},
   115  		}, inheritedSelectors)
   116  
   117  	// expectedPBSelectors should combine the inherited selectors and
   118  	// the additional selectors.
   119  	expectedPBSelectors := []*pb.Selector{
   120  		{
   121  			Label:  "label1",
   122  			Target: "value1",
   123  			Op:     pb.Selector_Eq,
   124  		},
   125  		{
   126  			Label:  "label2",
   127  			Target: "value2",
   128  			Op:     pb.Selector_Neq,
   129  		},
   130  		{
   131  			Label:  "label3",
   132  			Target: "value3",
   133  			Op:     pb.Selector_Eq,
   134  		},
   135  		{
   136  			Label:  "label4",
   137  			Target: "value4",
   138  			Op:     pb.Selector_Neq,
   139  		},
   140  	}
   141  	helper.ServerMasterClient.EXPECT().ScheduleTask(gomock.Any(),
   142  		&pb.ScheduleTaskRequest{
   143  			TaskId: "worker-1",
   144  			Resources: resModel.ToResourceRequirement(
   145  				"job-1", "/local/resource-1", "/local/resource-2"),
   146  			Selectors: expectedPBSelectors,
   147  		}).Return(
   148  		&pb.ScheduleTaskResponse{
   149  			ExecutorId:   "executor-1",
   150  			ExecutorAddr: "1.1.1.1:1234",
   151  		}, nil).Times(1)
   152  	executorCli := client.NewMockExecutorClient(gomock.NewController(t))
   153  	helper.ExecutorGroup.AddClient("executor-1", executorCli)
   154  	projectInfo := tenant.NewProjectInfo("tenant-1", "project-1")
   155  	executorCli.EXPECT().DispatchTask(gomock.Any(),
   156  		&client.DispatchTaskArgs{
   157  			ProjectInfo:  projectInfo,
   158  			WorkerID:     "worker-1",
   159  			MasterID:     "job-1",
   160  			WorkerType:   int64(1),
   161  			WorkerConfig: []byte("sample-config"),
   162  			WorkerEpoch:  int64(1),
   163  		}, gomock.Any()).
   164  		Return(nil).
   165  		Times(1)
   166  
   167  	err := helper.Creator.CreateWorker(
   168  		context.Background(),
   169  		projectInfo,
   170  		frameModel.WorkerType(1),
   171  		"worker-1",
   172  		[]byte("sample-config"),
   173  		CreateWorkerWithSelectors(additionalSelectors...),
   174  		CreateWorkerWithResourceRequirements("/local/resource-1", "/local/resource-2"))
   175  	require.NoError(t, err)
   176  }