github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpollmux/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 netpollmux

import (
	"context"
	"time"

	"github.com/cloudwego/netpoll"

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

func newMockRPCInfo() rpcinfo.RPCInfo {
	method := "method"
	c := rpcinfo.NewEndpointInfo("", method, nil, nil)
	s := rpcinfo.NewEndpointInfo("", method, nil, nil)
	ink := rpcinfo.NewInvocation("", method)
	cfg := rpcinfo.NewRPCConfig()
	ri := rpcinfo.NewRPCInfo(c, s, ink, cfg, 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 _ 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()
	}
}

var _ netpoll.Connection = &MockNetpollConn{}

// MockNetpollConn implements netpoll.Connection.
type MockNetpollConn struct {
	mocks.Conn
	AddCloseCallbackFunc func(c netpoll.CloseCallback) (e error)
	IsActiveFunc         func() (r bool)
	ReaderFunc           func() (r netpoll.Reader)
	WriterFunc           func() (r netpoll.Writer)
	SetIdleTimeoutFunc   func(timeout time.Duration) (e error)
	SetOnRequestFunc     func(on netpoll.OnRequest) (e error)
	SetReadTimeoutFunc   func(timeout time.Duration) (e error)
	SetWriteTimeoutFunc  func(timeout time.Duration) (e error)
}

// AddCloseCallback implements the netpoll.Connection interface.
func (m *MockNetpollConn) AddCloseCallback(c netpoll.CloseCallback) (e error) {
	if m.AddCloseCallbackFunc != nil {
		return m.AddCloseCallbackFunc(c)
	}
	return
}

// IsActive implements the netpoll.Connection interface.
func (m *MockNetpollConn) IsActive() (r bool) {
	if m.IsActiveFunc != nil {
		return m.IsActiveFunc()
	}
	return
}

// Reader implements the netpoll.Connection interface.
func (m *MockNetpollConn) Reader() (r netpoll.Reader) {
	if m.ReaderFunc != nil {
		return m.ReaderFunc()
	}
	return
}

// Writer implements the netpoll.Connection interface.
func (m *MockNetpollConn) Writer() (r netpoll.Writer) {
	if m.WriterFunc != nil {
		return m.WriterFunc()
	}
	return
}

// SetIdleTimeout implements the netpoll.Connection interface.
func (m *MockNetpollConn) SetIdleTimeout(timeout time.Duration) (e error) {
	if m.SetIdleTimeoutFunc != nil {
		return m.SetIdleTimeoutFunc(timeout)
	}
	return
}

// SetOnRequest implements the netpoll.Connection interface.
func (m *MockNetpollConn) SetOnRequest(on netpoll.OnRequest) (e error) {
	if m.SetOnRequestFunc != nil {
		return m.SetOnRequestFunc(on)
	}
	return
}

// SetReadTimeout implements the netpoll.Connection interface.
func (m *MockNetpollConn) SetReadTimeout(timeout time.Duration) (e error) {
	if m.SetReadTimeoutFunc != nil {
		return m.SetReadTimeoutFunc(timeout)
	}
	return
}

// SetWriteTimeout implements the netpoll.Connection interface.
func (m *MockNetpollConn) SetWriteTimeout(timeout time.Duration) (e error) {
	if m.SetWriteTimeoutFunc != nil {
		return m.SetWriteTimeoutFunc(timeout)
	}
	return
}

// MockNetpollWriter implements netpoll.Writer
type MockNetpollWriter struct {
	FlushFunc       func() (err error)
	WriteDirectFunc func(p []byte, remainCap int) error
}

var _ netpoll.Writer = &MockNetpollWriter{}

// Flush implements the netpoll.Writer interface.
func (m *MockNetpollWriter) Flush() (err error) {
	if m.FlushFunc != nil {
		return m.FlushFunc()
	}
	return
}

// Malloc implements the netpoll.Writer interface.
func (m *MockNetpollWriter) Malloc(n int) (buf []byte, err error) {
	return
}

// MallocLen implements the netpoll.Writer interface.
func (m *MockNetpollWriter) MallocLen() (length int) {
	return
}

// MallocAck implements the netpoll.Writer interface.
func (m *MockNetpollWriter) MallocAck(n int) (err error) {
	return
}

// Append implements the netpoll.Writer interface.
func (m *MockNetpollWriter) Append(w netpoll.Writer) (err error) {
	return
}

// Write implements the netpoll.Writer interface.
func (m *MockNetpollWriter) Write(b []byte) (n int, err error) {
	return
}

// WriteString implements the netpoll.Writer interface.
func (m *MockNetpollWriter) WriteString(s string) (n int, err error) {
	return
}

// WriteBinary implements the netpoll.Writer interface.
func (m *MockNetpollWriter) WriteBinary(b []byte) (n int, err error) {
	return
}

// WriteDirect implements the netpoll.Writer interface.
func (m *MockNetpollWriter) WriteDirect(p []byte, remainCap int) (err error) {
	if m.WriteDirectFunc != nil {
		return m.WriteDirectFunc(p, remainCap)
	}
	return
}

// WriteByte implements the netpoll.Writer interface.
func (m *MockNetpollWriter) WriteByte(b byte) (err error) {
	return
}

// MockNetpollReader implements netpoll.Reader
type MockNetpollReader struct {
	ReleaseFunc func() (err error)
}

// Next implements the netpoll.Reader interface.
func (m *MockNetpollReader) Next(n int) (p []byte, err error) {
	return
}

// Peek implements the netpoll.Reader interface.
func (m *MockNetpollReader) Peek(n int) (buf []byte, err error) {
	return
}

// Until implements the netpoll.Reader interface.
func (m *MockNetpollReader) Until(b byte) (buf []byte, err error) {
	return
}

// Skip implements the netpoll.Reader interface.
func (m *MockNetpollReader) Skip(n int) (err error) {
	return
}

// Release implements the netpoll.Reader interface.
func (m *MockNetpollReader) Release() (err error) {
	if m.ReleaseFunc != nil {
		return m.ReleaseFunc()
	}
	return
}

// Slice implements the netpoll.Reader interface.
func (m *MockNetpollReader) Slice(n int) (r netpoll.Reader, err error) {
	return
}

// Len implements the netpoll.Reader interface.
func (m *MockNetpollReader) Len() (length int) {
	return
}

// Read implements the netpoll.Reader interface.
func (m *MockNetpollReader) Read(b []byte) (n int, err error) {
	return
}

// ReadString implements the netpoll.Reader interface.
func (m *MockNetpollReader) ReadString(n int) (s string, err error) {
	return
}

// ReadBinary implements the netpoll.Reader interface.
func (m *MockNetpollReader) ReadBinary(n int) (p []byte, err error) {
	return
}

// ReadByte implements the netpoll.Reader interface.
func (m *MockNetpollReader) ReadByte() (b byte, err error) {
	return
}