github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/mocks_test.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package trans

import (
	"context"
	"net"

	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

func newMockRPCInfo() rpcinfo.RPCInfo {
	c := rpcinfo.NewEndpointInfo("", "", nil, nil)
	s := rpcinfo.NewEndpointInfo("", "", nil, nil)
	ink := rpcinfo.NewInvocation("", "")
	ri := rpcinfo.NewRPCInfo(c, s, ink, nil, rpcinfo.NewRPCStats())
	return ri
}

var _ remote.Codec = &MockCodec{}

type MockCodec struct {
	EncodeFunc func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error
	DecodeFunc func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error
	NameFunc   func() string
}

func (m *MockCodec) Encode(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
	if m.EncodeFunc != nil {
		return m.EncodeFunc(ctx, msg, out)
	}
	return nil
}

func (m *MockCodec) Decode(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
	if m.DecodeFunc != nil {
		return m.DecodeFunc(ctx, msg, in)
	}
	return nil
}

func (m *MockCodec) Name() string {
	if m.NameFunc != nil {
		return m.NameFunc()
	}
	return ""
}

var _ Extension = &MockExtension{}

type MockExtension struct {
	SetReadTimeoutFunc     func(ctx context.Context, conn net.Conn, cfg rpcinfo.RPCConfig, role remote.RPCRole)
	NewWriteByteBufferFunc func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer
	NewReadByteBufferFunc  func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer
	ReleaseBufferFunc      func(remote.ByteBuffer, error) error
	IsTimeoutErrFunc       func(error) bool
	IsRemoteClosedErrFunc  func(error) bool
}

func (m *MockExtension) SetReadTimeout(ctx context.Context, conn net.Conn, cfg rpcinfo.RPCConfig, role remote.RPCRole) {
	if m.SetReadTimeoutFunc != nil {
		m.SetReadTimeoutFunc(ctx, conn, cfg, role)
		return
	}
}

func (m *MockExtension) NewWriteByteBuffer(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
	if m.NewWriteByteBufferFunc != nil {
		return m.NewWriteByteBufferFunc(ctx, conn, msg)
	}
	return nil
}

func (m *MockExtension) NewReadByteBuffer(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
	if m.NewReadByteBufferFunc != nil {
		return m.NewReadByteBufferFunc(ctx, conn, msg)
	}
	return nil
}

func (m *MockExtension) ReleaseBuffer(buf remote.ByteBuffer, err error) error {
	if m.ReleaseBufferFunc != nil {
		return m.ReleaseBufferFunc(buf, err)
	}
	return nil
}

func (m *MockExtension) IsTimeoutErr(err error) bool {
	if m.IsTimeoutErrFunc != nil {
		return m.IsTimeoutErrFunc(err)
	}
	return false
}

func (m *MockExtension) IsRemoteClosedErr(err error) bool {
	if m.IsRemoteClosedErrFunc != nil {
		return m.IsRemoteClosedErrFunc(err)
	}
	return false
}

var _ remote.Message = &MockMessage{}

type MockMessage struct {
	RPCInfoFunc         func() rpcinfo.RPCInfo
	ServiceInfoFunc     func() *serviceinfo.ServiceInfo
	SetServiceInfoFunc  func(svcName, methodName string) (*serviceinfo.ServiceInfo, error)
	DataFunc            func() interface{}
	NewDataFunc         func(method string) (ok bool)
	MessageTypeFunc     func() remote.MessageType
	SetMessageTypeFunc  func(remote.MessageType)
	RPCRoleFunc         func() remote.RPCRole
	PayloadLenFunc      func() int
	SetPayloadLenFunc   func(size int)
	TransInfoFunc       func() remote.TransInfo
	TagsFunc            func() map[string]interface{}
	ProtocolInfoFunc    func() remote.ProtocolInfo
	SetProtocolInfoFunc func(remote.ProtocolInfo)
	PayloadCodecFunc    func() remote.PayloadCodec
	SetPayloadCodecFunc func(pc remote.PayloadCodec)
	RecycleFunc         func()
}

func (m *MockMessage) RPCInfo() rpcinfo.RPCInfo {
	if m.RPCInfoFunc != nil {
		return m.RPCInfoFunc()
	}
	return nil
}

func (m *MockMessage) ServiceInfo() (si *serviceinfo.ServiceInfo) {
	if m.ServiceInfoFunc != nil {
		return m.ServiceInfoFunc()
	}
	return
}

func (m *MockMessage) SpecifyServiceInfo(svcName, methodName string) (si *serviceinfo.ServiceInfo, err error) {
	if m.SetServiceInfoFunc != nil {
		return m.SetServiceInfoFunc(svcName, methodName)
	}
	return nil, nil
}

func (m *MockMessage) Data() interface{} {
	if m.DataFunc != nil {
		return m.DataFunc()
	}
	return nil
}

func (m *MockMessage) NewData(method string) (ok bool) {
	if m.NewDataFunc != nil {
		return m.NewDataFunc(method)
	}
	return false
}

func (m *MockMessage) MessageType() (mt remote.MessageType) {
	if m.MessageTypeFunc != nil {
		return m.MessageTypeFunc()
	}
	return
}

func (m *MockMessage) SetMessageType(mt remote.MessageType) {
	if m.SetMessageTypeFunc != nil {
		m.SetMessageTypeFunc(mt)
	}
}

func (m *MockMessage) RPCRole() (r remote.RPCRole) {
	if m.RPCRoleFunc != nil {
		return m.RPCRoleFunc()
	}
	return
}

func (m *MockMessage) PayloadLen() int {
	if m.PayloadLenFunc != nil {
		return m.PayloadLenFunc()
	}
	return 0
}

func (m *MockMessage) SetPayloadLen(size int) {
	if m.SetPayloadLenFunc != nil {
		m.SetPayloadLenFunc(size)
	}
}

func (m *MockMessage) TransInfo() remote.TransInfo {
	if m.TransInfoFunc != nil {
		return m.TransInfoFunc()
	}
	return nil
}

func (m *MockMessage) Tags() map[string]interface{} {
	if m.TagsFunc != nil {
		return m.TagsFunc()
	}
	return nil
}

func (m *MockMessage) ProtocolInfo() (pi remote.ProtocolInfo) {
	if m.ProtocolInfoFunc != nil {
		return m.ProtocolInfoFunc()
	}
	return
}

func (m *MockMessage) SetProtocolInfo(pi remote.ProtocolInfo) {
	if m.SetProtocolInfoFunc != nil {
		m.SetProtocolInfoFunc(pi)
	}
}

func (m *MockMessage) PayloadCodec() remote.PayloadCodec {
	if m.PayloadCodecFunc != nil {
		return m.PayloadCodecFunc()
	}
	return nil
}

func (m *MockMessage) SetPayloadCodec(pc remote.PayloadCodec) {
	if m.SetPayloadCodecFunc != nil {
		m.SetPayloadCodecFunc(pc)
	}
}

func (m *MockMessage) Recycle() {
	if m.RecycleFunc != nil {
		m.RecycleFunc()
	}
}