github.com/cloudwego/kitex@v0.9.0/internal/mocks/serviceinfo.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 mocks

import (
	"context"
	"errors"
	"fmt"

	"github.com/apache/thrift/lib/go/thrift"

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

// method name of mock
const (
	MockServiceName            = "MockService"
	MockService2Name           = "MockService2"
	MockService3Name           = "MockService3"
	MockMethod          string = "mock"
	Mock2Method         string = "mock2"
	MockExceptionMethod string = "mockException"
	MockErrorMethod     string = "mockError"
	MockOnewayMethod    string = "mockOneway"
)

// ServiceInfo return mock serviceInfo
func ServiceInfo() *serviceinfo.ServiceInfo {
	return myServiceServiceInfo
}

var myServiceServiceInfo = newServiceInfo()

func newServiceInfo() *serviceinfo.ServiceInfo {
	methods := map[string]serviceinfo.MethodInfo{
		"mock":          serviceinfo.NewMethodInfo(mockHandler, NewMockArgs, NewMockResult, false),
		"mockException": serviceinfo.NewMethodInfo(mockExceptionHandler, NewMockArgs, newMockExceptionResult, false),
		"mockError":     serviceinfo.NewMethodInfo(mockErrorHandler, NewMockArgs, NewMockResult, false),
		"mockOneway":    serviceinfo.NewMethodInfo(mockOnewayHandler, NewMockArgs, nil, true),
		"mockStreaming": serviceinfo.NewMethodInfo(
			nil, nil, nil, false,
			serviceinfo.WithStreamingMode(serviceinfo.StreamingBidirectional),
		),
	}

	svcInfo := &serviceinfo.ServiceInfo{
		ServiceName: MockServiceName,
		Methods:     methods,
		Extra: map[string]interface{}{
			"PackageName": "mock",
		},
	}
	return svcInfo
}

// Service2Info return mock serviceInfo
func Service2Info() *serviceinfo.ServiceInfo {
	return myServiceService2Info
}

var myServiceService2Info = newService2Info()

func newService2Info() *serviceinfo.ServiceInfo {
	methods := map[string]serviceinfo.MethodInfo{
		"mock2": serviceinfo.NewMethodInfo(mock2Handler, NewMockArgs, NewMockResult, false),
	}

	svcInfo := &serviceinfo.ServiceInfo{
		ServiceName: MockService2Name,
		Methods:     methods,
		Extra: map[string]interface{}{
			"PackageName": "mock2",
		},
	}
	return svcInfo
}

// Service3Info return mock serviceInfo
func Service3Info() *serviceinfo.ServiceInfo {
	return myServiceService3Info
}

var myServiceService3Info = newService3Info()

func newService3Info() *serviceinfo.ServiceInfo {
	methods := map[string]serviceinfo.MethodInfo{
		"mock": serviceinfo.NewMethodInfo(mockHandler, NewMockArgs, NewMockResult, false),
	}

	svcInfo := &serviceinfo.ServiceInfo{
		ServiceName: MockService3Name,
		Methods:     methods,
		Extra: map[string]interface{}{
			"PackageName": "mock",
		},
	}
	return svcInfo
}

func mockHandler(ctx context.Context, handler, args, result interface{}) error {
	a := args.(*myServiceMockArgs)
	r := result.(*myServiceMockResult)
	reply, err := handler.(MyService).Mock(ctx, a.Req)
	if err != nil {
		return err
	}
	r.Success = reply
	return nil
}

func mock2Handler(ctx context.Context, handler, args, result interface{}) error {
	a := args.(*myServiceMockArgs)
	r := result.(*myServiceMockResult)
	reply, err := handler.(MyService).Mock(ctx, a.Req)
	if err != nil {
		return err
	}
	r.Success = reply
	return nil
}

func NewMockArgs() interface{} {
	return &myServiceMockArgs{}
}

func NewMockResult() interface{} {
	return &myServiceMockResult{}
}

func mockExceptionHandler(ctx context.Context, handler, args, result interface{}) error {
	a := args.(*myServiceMockArgs)
	r := result.(*myServiceMockExceptionResult)
	reply, err := handler.(MyService).MockException(ctx, a.Req)
	if err != nil {
		switch v := err.(type) {
		case *MyException:
			r.MyException = v
		default:
			return err
		}
	} else {
		r.Success = reply
	}
	return nil
}

func mockErrorHandler(ctx context.Context, handler, args, result interface{}) error {
	a := args.(*myServiceMockArgs)
	r := result.(*myServiceMockResult)
	reply, err := handler.(MyService).MockError(ctx, a.Req)
	if err != nil {
		return err
	}
	r.Success = reply
	return nil
}

func newMockExceptionResult() interface{} {
	return &myServiceMockExceptionResult{}
}

func mockOnewayHandler(ctx context.Context, handler, args, result interface{}) error {
	a := args.(*myServiceMockArgs)
	err := handler.(MyService).MockOneway(ctx, a.Req)
	if err != nil {
		return err
	}
	return nil
}

// MyService .
type MyService interface {
	Mock(ctx context.Context, req *MyRequest) (r *MyResponse, err error)
	MockException(ctx context.Context, req *MyRequest) (r *MyResponse, err error)
	MockError(ctx context.Context, req *MyRequest) (r *MyResponse, err error)
	MockOneway(ctx context.Context, req *MyRequest) (err error)
	Mock2(ctx context.Context, req *MyRequest) (r *MyResponse, err error)
}

type myServiceMockArgs struct {
	Req *MyRequest `thrift:"req,1" json:"req"`
}

func (p *myServiceMockArgs) Read(iprot thrift.TProtocol) error {
	return nil
}

func (p *myServiceMockArgs) Write(oprot thrift.TProtocol) error {
	return nil
}

// MyRequest .
type MyRequest struct {
	Name string `thrift:"Name,1" json:"Name"`
}

type myServiceMockResult struct {
	Success *MyResponse `thrift:"success,0" json:"success,omitempty"`
}

func (p *myServiceMockResult) Read(iprot thrift.TProtocol) error {
	return nil
}

func (p *myServiceMockResult) Write(oprot thrift.TProtocol) error {
	return nil
}

// MyResponse .
type MyResponse struct {
	Name string `thrift:"Name,1" json:"Name"`
}

type myServiceMockExceptionResult struct {
	Success     *MyResponse  `thrift:"success,0" json:"success,omitempty"`
	MyException *MyException `thrift:"stException,1" json:"stException,omitempty"`
}

func (p *myServiceMockExceptionResult) Read(iprot thrift.TProtocol) error {
	return nil
}

func (p *myServiceMockExceptionResult) Write(oprot thrift.TProtocol) error {
	return nil
}

// MyException .
type MyException struct {
	Message string `thrift:"message,1" json:"message"`
}

func (p *MyException) Error() string {
	return fmt.Sprintf("MyException(%+v)", *p)
}

// MyServiceHandler .
func MyServiceHandler() interface{} {
	return &myServiceHandler{}
}

// MockFuncHandler .
func MockFuncHandler(mf func(ctx context.Context, req *MyRequest) (r *MyResponse, err error)) interface{} {
	return &myServiceHandler{mf}
}

type myServiceHandler struct {
	mockFunc func(ctx context.Context, req *MyRequest) (r *MyResponse, err error)
}

func (h *myServiceHandler) Mock(ctx context.Context, req *MyRequest) (r *MyResponse, err error) {
	if h.mockFunc != nil {
		return h.mockFunc(ctx, req)
	}
	return &MyResponse{Name: MockMethod}, nil
}

func (h *myServiceHandler) Mock2(ctx context.Context, req *MyRequest) (r *MyResponse, err error) {
	if h.mockFunc != nil {
		return h.mockFunc(ctx, req)
	}
	return &MyResponse{Name: Mock2Method}, nil
}

func (h *myServiceHandler) MockException(ctx context.Context, req *MyRequest) (r *MyResponse, err error) {
	return &MyResponse{Name: MockExceptionMethod}, nil
}

func (h *myServiceHandler) MockError(ctx context.Context, req *MyRequest) (r *MyResponse, err error) {
	return nil, errors.New(MockErrorMethod)
}

func (h *myServiceHandler) MockOneway(ctx context.Context, req *MyRequest) (err error) {
	return nil
}