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

import (
	"context"
	"time"

	"github.com/cloudwego/netpoll"

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

var _ netpoll.Connection = &MockNetpollConn{}

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
}

func (m *MockCodec) Name() string {
	return "Mock"
}

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

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

// 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
}

// SetReadTimeout 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) (p []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
}