github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/framework/registry/factory_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 registry
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	"github.com/pingcap/tiflow/engine/framework"
    21  	frameModel "github.com/pingcap/tiflow/engine/framework/model"
    22  	dcontext "github.com/pingcap/tiflow/engine/pkg/context"
    23  	"github.com/pingcap/tiflow/engine/pkg/deps"
    24  	"github.com/pingcap/tiflow/engine/pkg/externalresource/broker"
    25  	"github.com/pingcap/tiflow/engine/pkg/meta/mock"
    26  	metaModel "github.com/pingcap/tiflow/engine/pkg/meta/model"
    27  	pkgOrm "github.com/pingcap/tiflow/engine/pkg/orm"
    28  	"github.com/pingcap/tiflow/engine/pkg/p2p"
    29  	"github.com/stretchr/testify/require"
    30  	"go.uber.org/dig"
    31  )
    32  
    33  type paramList struct {
    34  	dig.Out
    35  
    36  	MessageHandlerManager p2p.MessageHandlerManager
    37  	MessageSender         p2p.MessageSender
    38  	FrameMetaClient       pkgOrm.Client
    39  	BusinessClientConn    metaModel.ClientConn
    40  	ResourceBroker        broker.Broker
    41  }
    42  
    43  func makeCtxWithMockDeps(t *testing.T) (*dcontext.Context, context.CancelFunc) {
    44  	dp := deps.NewDeps()
    45  	cli, err := pkgOrm.NewMockClient()
    46  	require.NoError(t, err)
    47  	broker := broker.NewBrokerForTesting("executor-1")
    48  	err = dp.Provide(func() paramList {
    49  		return paramList{
    50  			MessageHandlerManager: p2p.NewMockMessageHandlerManager(),
    51  			MessageSender:         p2p.NewMockMessageSender(),
    52  			FrameMetaClient:       cli,
    53  			BusinessClientConn:    mock.NewMockClientConn(),
    54  			ResourceBroker:        broker,
    55  		}
    56  	})
    57  	require.NoError(t, err)
    58  
    59  	cancelFn := func() {
    60  		broker.Close()
    61  	}
    62  	return dcontext.Background().WithDeps(dp), cancelFn
    63  }
    64  
    65  type fakeWorkerConfig struct {
    66  	TargetTick int `json:"target-tick"`
    67  }
    68  
    69  type fakeWorker struct {
    70  	framework.WorkerImpl
    71  	framework.BaseWorker
    72  }
    73  
    74  func newFakeWorker(_ *dcontext.Context, _ frameModel.WorkerID, _ frameModel.MasterID, _ *fakeWorkerConfig) framework.WorkerImpl {
    75  	return &fakeWorker{}
    76  }
    77  
    78  func TestNewSimpleWorkerFactory(t *testing.T) {
    79  	fac := NewSimpleWorkerFactory(newFakeWorker)
    80  	config, err := fac.DeserializeConfig([]byte(`{"target-tick":100}`))
    81  	require.NoError(t, err)
    82  	require.Equal(t, &fakeWorkerConfig{TargetTick: 100}, config)
    83  
    84  	ctx, cancel := makeCtxWithMockDeps(t)
    85  	defer cancel()
    86  
    87  	metaCli, err := ctx.Deps().Construct(
    88  		func(cli pkgOrm.Client) (pkgOrm.Client, error) {
    89  			return cli, nil
    90  		},
    91  	)
    92  	require.NoError(t, err)
    93  	defer metaCli.(pkgOrm.Client).Close()
    94  	newWorker, err := fac.NewWorkerImpl(ctx, "my-worker", "my-master", &fakeWorkerConfig{TargetTick: 100})
    95  	require.NoError(t, err)
    96  	require.IsType(t, &fakeWorker{}, newWorker)
    97  }
    98  
    99  func TestDeserializeConfigError(t *testing.T) {
   100  	fac := NewSimpleWorkerFactory(newFakeWorker)
   101  	_, err := fac.DeserializeConfig([]byte(`{target-tick:100}`))
   102  	require.Error(t, err)
   103  	require.False(t, fac.IsRetryableError(err))
   104  }