github.com/matrixorigin/matrixone@v1.2.0/pkg/cnservice/server_test.go (about)

     1  // Copyright 2023 Matrix Origin
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package cnservice
    16  
    17  import (
    18  	"context"
    19  	"sync"
    20  	"testing"
    21  
    22  	"github.com/golang/mock/gomock"
    23  	"github.com/matrixorigin/matrixone/pkg/common/morpc"
    24  	"github.com/matrixorigin/matrixone/pkg/common/morpc/mock_morpc"
    25  	"github.com/matrixorigin/matrixone/pkg/defines"
    26  	"github.com/matrixorigin/matrixone/pkg/fileservice"
    27  	"github.com/matrixorigin/matrixone/pkg/lockservice"
    28  	"github.com/matrixorigin/matrixone/pkg/logservice"
    29  	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
    30  	"github.com/matrixorigin/matrixone/pkg/pb/pipeline"
    31  	qclient "github.com/matrixorigin/matrixone/pkg/queryservice/client"
    32  	"github.com/matrixorigin/matrixone/pkg/txn/client"
    33  	"github.com/matrixorigin/matrixone/pkg/udf"
    34  	"github.com/matrixorigin/matrixone/pkg/util/address"
    35  	"github.com/matrixorigin/matrixone/pkg/vm/engine"
    36  	"github.com/stretchr/testify/require"
    37  )
    38  
    39  func Test_InitServer(t *testing.T) {
    40  	ctrl := gomock.NewController(t)
    41  	defer ctrl.Finish()
    42  
    43  	cfg := &Config{
    44  		UUID:     "dd1dccb4-4d3c-41f8-b482-5251dc7a41bf",
    45  		PortBase: 18000,
    46  	}
    47  
    48  	srv := &service{
    49  		metadata: metadata.CNStore{
    50  			UUID: cfg.UUID,
    51  		},
    52  		cfg: cfg,
    53  		responsePool: &sync.Pool{
    54  			New: func() any {
    55  				return &pipeline.Message{}
    56  			},
    57  		},
    58  		addressMgr: address.NewAddressManager(cfg.ServiceHost, cfg.PortBase),
    59  	}
    60  	srv.addressMgr.Register(0)
    61  
    62  	WithTaskStorageFactory(nil)(srv)
    63  	handler := func(
    64  		ctx context.Context,
    65  		cnAddr string,
    66  		message morpc.Message,
    67  		cs morpc.ClientSession,
    68  		engine engine.Engine,
    69  		fs fileservice.FileService,
    70  		lockService lockservice.LockService,
    71  		queryClient qclient.QueryClient,
    72  		hakeeper logservice.CNHAKeeperClient,
    73  		udfService udf.Service,
    74  		cli client.TxnClient,
    75  		aicm *defines.AutoIncrCacheManager,
    76  		mAcquirer func() morpc.Message) error {
    77  		return nil
    78  	}
    79  	WithMessageHandle(handler)(srv)
    80  
    81  	require.Equal(t, srv.ID(), cfg.UUID)
    82  	require.Equal(t, srv.SQLAddress(), cfg.SQLAddress)
    83  
    84  	msg := &pipeline.Message{}
    85  
    86  	srv.releaseMessage(msg)
    87  	message := srv.acquireMessage()
    88  	require.Equal(t, message.(*pipeline.Message).Sid, msg.Sid)
    89  
    90  	var err error
    91  	ctx := context.TODO()
    92  	session := mock_morpc.NewMockClientSession(ctrl)
    93  	msg.Cmd = pipeline.Method_PipelineMessage
    94  	session.EXPECT().CreateCache(ctx, uint64(0)).Return(&testMessageCache{}, nil).Times(2)
    95  
    96  	msg.Sid = pipeline.Status_WaitingNext
    97  	err = srv.handleRequest(
    98  		ctx,
    99  		morpc.RPCMessage{
   100  			Ctx:     ctx,
   101  			Cancel:  func() {},
   102  			Message: msg,
   103  		},
   104  		0,
   105  		session,
   106  	)
   107  	require.Nil(t, err)
   108  
   109  	msg.Sid = pipeline.Status_Last
   110  	err = srv.handleRequest(
   111  		ctx,
   112  		morpc.RPCMessage{
   113  			Ctx:     ctx,
   114  			Cancel:  func() {},
   115  			Message: msg,
   116  		},
   117  		0,
   118  		session,
   119  	)
   120  	require.Nil(t, err)
   121  }
   122  
   123  type testMessageCache struct {
   124  	cache []morpc.Message
   125  }
   126  
   127  func (c *testMessageCache) Add(val morpc.Message) error {
   128  	c.cache = append(c.cache, val)
   129  	return nil
   130  }
   131  
   132  func (c *testMessageCache) Len() (int, error) {
   133  	return len(c.cache), nil
   134  }
   135  
   136  func (c *testMessageCache) Pop() (morpc.Message, bool, error) {
   137  	if len(c.cache) == 0 {
   138  		return nil, false, nil
   139  	}
   140  	ret := c.cache[0]
   141  	c.cache = c.cache[1:]
   142  	return ret, true, nil
   143  }
   144  
   145  func (c *testMessageCache) Close() {
   146  }