github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/servermaster/scheduler/scheduler_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 scheduler
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/pingcap/tiflow/engine/model"
    21  	resModel "github.com/pingcap/tiflow/engine/pkg/externalresource/model"
    22  	schedModel "github.com/pingcap/tiflow/engine/servermaster/scheduler/model"
    23  	"github.com/pingcap/tiflow/pkg/errors"
    24  	"github.com/pingcap/tiflow/pkg/label"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  func getMockDataForScheduler() *mockExecutorInfoProvider {
    29  	return &mockExecutorInfoProvider{
    30  		infos: map[model.ExecutorID]schedModel.ExecutorInfo{
    31  			"executor-1": {
    32  				ID: "executor-1",
    33  				Labels: label.Set{
    34  					"type":     "test-type",
    35  					"function": "test-function",
    36  				},
    37  			},
    38  			"executor-2": {ID: "executor-2"},
    39  			"executor-3": {ID: "executor-3"},
    40  		},
    41  	}
    42  }
    43  
    44  func getMockResourceConstraintForScheduler() PlacementConstrainer {
    45  	return &MockPlacementConstrainer{ResourceList: map[resModel.ResourceKey]model.ExecutorID{
    46  		{JobID: "fakeJob", ID: "resource-1"}: "executor-1",
    47  		{JobID: "fakeJob", ID: "resource-2"}: "executor-2",
    48  		{JobID: "fakeJob", ID: "resource-3"}: "executor-3",
    49  		{JobID: "fakeJob", ID: "resource-4"}: "", // no constraint
    50  	}}
    51  }
    52  
    53  func TestSchedulerByConstraint(t *testing.T) {
    54  	sched := NewScheduler(
    55  		getMockDataForScheduler(),
    56  		getMockResourceConstraintForScheduler())
    57  
    58  	resp, err := sched.ScheduleTask(context.Background(), &schedModel.SchedulerRequest{
    59  		ExternalResources: []resModel.ResourceKey{{JobID: "fakeJob", ID: "resource-2"}},
    60  	})
    61  	require.NoError(t, err)
    62  	require.Equal(t, &schedModel.SchedulerResponse{ExecutorID: "executor-2"}, resp)
    63  }
    64  
    65  func TestSchedulerNoResourceConstraint(t *testing.T) {
    66  	sched := NewScheduler(
    67  		getMockDataForScheduler(),
    68  		getMockResourceConstraintForScheduler())
    69  
    70  	resp, err := sched.ScheduleTask(context.Background(), &schedModel.SchedulerRequest{
    71  		// resource-4 has no constraint
    72  		ExternalResources: []resModel.ResourceKey{{JobID: "fakeJob", ID: "resource-4"}},
    73  		Selectors: []*label.Selector{
    74  			{
    75  				Key:    "type",
    76  				Target: "test-type",
    77  				Op:     "eq",
    78  			},
    79  			{
    80  				Key:    "function",
    81  				Target: "test-function",
    82  				Op:     "eq",
    83  			},
    84  		},
    85  	})
    86  	require.NoError(t, err)
    87  	require.Equal(t, &schedModel.SchedulerResponse{ExecutorID: "executor-1"}, resp)
    88  }
    89  
    90  func TestSchedulerResourceNotFound(t *testing.T) {
    91  	sched := NewScheduler(
    92  		getMockDataForScheduler(),
    93  		getMockResourceConstraintForScheduler())
    94  
    95  	_, err := sched.ScheduleTask(context.Background(), &schedModel.SchedulerRequest{
    96  		// resource-blah DOES NOT exist
    97  		ExternalResources: []resModel.ResourceKey{{JobID: "fakeJob", ID: "resource-blah"}},
    98  	})
    99  	require.Error(t, err)
   100  	require.True(t, errors.Is(err, errors.ErrResourceDoesNotExist))
   101  }
   102  
   103  func TestSchedulerConstraintConflict(t *testing.T) {
   104  	sched := NewScheduler(
   105  		getMockDataForScheduler(),
   106  		getMockResourceConstraintForScheduler())
   107  
   108  	_, err := sched.ScheduleTask(context.Background(), &schedModel.SchedulerRequest{
   109  		ExternalResources: []resModel.ResourceKey{
   110  			{
   111  				JobID: "fakeJob",
   112  				ID:    "resource-1",
   113  			},
   114  			{
   115  				JobID: "fakeJob",
   116  				ID:    "resource-2",
   117  			},
   118  		},
   119  	})
   120  	require.Error(t, err)
   121  	require.True(t, errors.Is(err, errors.ErrResourceConflict))
   122  }
   123  
   124  func TestSchedulerNoQualifiedExecutor(t *testing.T) {
   125  	sched := NewScheduler(
   126  		&mockExecutorInfoProvider{infos: map[model.ExecutorID]schedModel.ExecutorInfo{}},
   127  		&MockPlacementConstrainer{},
   128  	)
   129  
   130  	_, err := sched.ScheduleTask(context.Background(), &schedModel.SchedulerRequest{})
   131  	require.ErrorIs(t, err, errors.ErrNoQualifiedExecutor)
   132  }