github.com/polarismesh/polaris@v1.17.8/service/batch/batch_test.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package batch
    19  
    20  import (
    21  	"context"
    22  	"fmt"
    23  	"sync"
    24  	"testing"
    25  
    26  	"github.com/golang/mock/gomock"
    27  	apiservice "github.com/polarismesh/specification/source/go/api/v1/service_manage"
    28  	. "github.com/smartystreets/goconvey/convey"
    29  	"github.com/stretchr/testify/assert"
    30  
    31  	"github.com/polarismesh/polaris/common/metrics"
    32  	"github.com/polarismesh/polaris/common/model"
    33  	"github.com/polarismesh/polaris/common/utils"
    34  	smock "github.com/polarismesh/polaris/store/mock"
    35  )
    36  
    37  func init() {
    38  	metrics.InitMetrics()
    39  }
    40  
    41  // TestNewBatchCtrlWithConfig 测试New
    42  func TestNewBatchCtrlWithConfig(t *testing.T) {
    43  	Convey("正常新建", t, func() {
    44  		ctrlConfig := &CtrlConfig{
    45  			Open:          true,
    46  			QueueSize:     1024,
    47  			WaitTime:      "20ms",
    48  			MaxBatchCount: 32,
    49  			Concurrency:   64,
    50  		}
    51  		config := &Config{
    52  			Register:   ctrlConfig,
    53  			Deregister: ctrlConfig,
    54  		}
    55  		bc, err := NewBatchCtrlWithConfig(nil, nil, config)
    56  		So(err, ShouldBeNil)
    57  		So(bc, ShouldNotBeNil)
    58  		So(bc.register, ShouldNotBeNil)
    59  		So(bc.deregister, ShouldNotBeNil)
    60  	})
    61  	Convey("可以关闭register和deregister的batch操作", t, func() {
    62  		bc, err := NewBatchCtrlWithConfig(nil, nil, nil)
    63  		So(err, ShouldBeNil)
    64  		So(bc, ShouldBeNil)
    65  
    66  		config := &Config{
    67  			Register:   &CtrlConfig{Open: false},
    68  			Deregister: &CtrlConfig{Open: false},
    69  		}
    70  		bc, err = NewBatchCtrlWithConfig(nil, nil, config)
    71  		So(err, ShouldBeNil)
    72  		So(bc, ShouldNotBeNil)
    73  		So(bc.register, ShouldBeNil)
    74  		So(bc.deregister, ShouldBeNil)
    75  	})
    76  }
    77  
    78  func newCreateInstanceController(t *testing.T) (*gomock.Controller, *Controller, *smock.MockStore, context.CancelFunc) {
    79  	ctl := gomock.NewController(t)
    80  	storage := smock.NewMockStore(ctl)
    81  	config := &Config{
    82  		Register: &CtrlConfig{
    83  			Open:          true,
    84  			QueueSize:     1024,
    85  			WaitTime:      "16ms",
    86  			MaxBatchCount: 8,
    87  			Concurrency:   4,
    88  		},
    89  	}
    90  	bc, err := NewBatchCtrlWithConfig(storage, nil, config)
    91  	if bc == nil || err != nil {
    92  		t.Fatalf("error: %+v", err)
    93  	}
    94  	ctx, cancel := context.WithCancel(context.Background())
    95  	bc.Start(ctx)
    96  	return ctl, bc, storage, cancel
    97  }
    98  
    99  func sendAsyncCreateInstance(bc *Controller, cnt int32) error {
   100  	var wg sync.WaitGroup
   101  	ch := make(chan error, cnt)
   102  	for i := int32(0); i < cnt; i++ {
   103  		wg.Add(1)
   104  		go func(index int32) {
   105  			defer wg.Done()
   106  			future := bc.AsyncCreateInstance(utils.NewUUID(), &apiservice.Instance{
   107  				Id:           utils.NewStringValue(fmt.Sprintf("%d", index)),
   108  				ServiceToken: utils.NewStringValue(fmt.Sprintf("%d", index)),
   109  			}, true)
   110  			if err := future.Wait(); err != nil {
   111  				fmt.Printf("%+v\n", err)
   112  				ch <- err
   113  			}
   114  		}(i)
   115  	}
   116  	wg.Wait()
   117  	select {
   118  	case err := <-ch:
   119  		if err != nil {
   120  			return err
   121  		}
   122  	default:
   123  		return nil
   124  	}
   125  	return nil
   126  }
   127  
   128  // TestAsyncCreateInstance test AsyncCreateInstance
   129  func TestAsyncCreateInstance(t *testing.T) {
   130  	t.Run("正常创建实例", func(t *testing.T) {
   131  		ctrl, bc, storage, cancel := newCreateInstanceController(t)
   132  		t.Cleanup(func() {
   133  			ctrl.Finish()
   134  			cancel()
   135  		})
   136  		mockSvc := &model.Service{ID: "1"}
   137  		totalIns := int32(100)
   138  		storage.EXPECT().BatchGetInstanceIsolate(gomock.Any()).Return(nil, nil).AnyTimes()
   139  		storage.EXPECT().GetSourceServiceToken(gomock.Any(), gomock.Any()).
   140  			Return(mockSvc, nil).AnyTimes()
   141  		storage.EXPECT().GetServiceByID(gomock.Any()).Return(mockSvc, nil).AnyTimes()
   142  		storage.EXPECT().BatchAddInstances(gomock.Any()).Return(nil).AnyTimes()
   143  		assert.NoError(t, sendAsyncCreateInstance(bc, totalIns))
   144  	})
   145  }
   146  
   147  // TestSendReply 测试reply
   148  func TestSendReply(t *testing.T) {
   149  	Convey("可以正常获取类型", t, func() {
   150  		sendReply(make([]*InstanceFuture, 0, 10), 1, nil)
   151  	})
   152  	Convey("可以正常获取类型2", t, func() {
   153  		sendReply(make(map[string]*InstanceFuture, 10), 1, nil)
   154  	})
   155  	Convey("其他类型不通过", t, func() {
   156  		sendReply("test string", 1, nil)
   157  	})
   158  }