github.com/cloudwego/kitex@v0.9.0/pkg/generic/binary_test/generic_init.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 test ...
package test

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"time"

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

	"github.com/cloudwego/kitex/client"
	"github.com/cloudwego/kitex/client/genericclient"
	kt "github.com/cloudwego/kitex/internal/mocks/thrift"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/server"
	"github.com/cloudwego/kitex/server/genericserver"
)

var (
	reqMsg  = "Hello Kitex"
	respMsg = "Hi, I am Kitex"
	errResp = "Test Error"
)

func newGenericClient(destService string, g generic.Generic, targetIPPort string, opts ...client.Option) genericclient.Client {
	opts = append(opts, client.WithHostPorts(targetIPPort))
	genericCli, _ := genericclient.NewClient(destService, g, opts...)
	return genericCli
}

func newGenericServer(g generic.Generic, addr net.Addr, handler generic.Service) server.Server {
	var opts []server.Option
	opts = append(opts, server.WithServiceAddr(addr), server.WithExitWaitTime(time.Microsecond*10))
	svr := genericserver.NewServer(handler, g, opts...)
	go func() {
		err := svr.Run()
		if err != nil {
			panic(err)
		}
	}()
	test.WaitServerStart(addr.String())
	return svr
}

// GenericServiceImpl ...
type GenericServiceImpl struct{}

// GenericCall ...
func (g *GenericServiceImpl) GenericCall(ctx context.Context, method string, request interface{}) (response interface{}, err error) {
	req := request.([]byte)
	fmt.Printf("Recv: %v\n", string(req[8+len(method):]))

	buf := genBinaryResp(method)
	return buf, nil
}

// GenericServiceErrorImpl ...
type GenericServiceErrorImpl struct{}

// GenericCall ...
func (g *GenericServiceErrorImpl) GenericCall(ctx context.Context, method string, request interface{}) (response interface{}, err error) {
	return response, errors.New(errResp)
}

// GenericServiceMockImpl ...
type GenericServiceMockImpl struct{}

// GenericCall ...
func (g *GenericServiceMockImpl) GenericCall(ctx context.Context, method string, request interface{}) (response interface{}, err error) {
	rc := utils.NewThriftMessageCodec()
	buf := request.([]byte)

	var args2 kt.MockTestArgs
	mth, seqID, err := rc.Decode(buf, &args2)
	if err != nil {
		return nil, err
	}

	if args2.Req.Msg != reqMsg {
		return nil, fmt.Errorf("msg is not %s", reqMsg)
	}

	resp := respMsg
	result := kt.NewMockTestResult()
	result.Success = &resp

	buf, err = rc.Encode(mth, thrift.REPLY, seqID, result)
	return buf, err
}

// NewMockServer normal server
func NewMockServer(handler kt.Mock, addr net.Addr, opts ...server.Option) server.Server {
	var options []server.Option
	opts = append(opts, server.WithServiceAddr(addr), server.WithExitWaitTime(time.Microsecond*10))
	options = append(options, opts...)

	svr := server.NewServer(options...)
	if err := svr.RegisterService(serviceInfo(), handler); err != nil {
		panic(err)
	}
	go func() {
		err := svr.Run()
		if err != nil {
			panic(err)
		}
	}()
	test.WaitServerStart(addr.String())
	return svr
}

func serviceInfo() *serviceinfo.ServiceInfo {
	destService := "Mock"
	handlerType := (*kt.Mock)(nil)
	methods := map[string]serviceinfo.MethodInfo{
		"Test": serviceinfo.NewMethodInfo(testHandler, newMockTestArgs, newMockTestResult, false),
	}
	svcInfo := &serviceinfo.ServiceInfo{
		ServiceName: destService,
		HandlerType: handlerType,
		Methods:     methods,
		Extra:       make(map[string]interface{}),
	}
	return svcInfo
}

func newMockTestArgs() interface{} {
	return kt.NewMockTestArgs()
}

func newMockTestResult() interface{} {
	return kt.NewMockTestResult()
}

func testHandler(ctx context.Context, handler, arg, result interface{}) error {
	realArg := arg.(*kt.MockTestArgs)
	realResult := result.(*kt.MockTestResult)
	success, err := handler.(kt.Mock).Test(ctx, realArg.Req)
	if err != nil {
		return err
	}
	realResult.Success = &success
	return nil
}

// MockImpl is used for test
type MockImpl struct{}

// Test ...
func (m *MockImpl) Test(ctx context.Context, req *kt.MockReq) (r string, err error) {
	if req.Msg != reqMsg {
		return "", fmt.Errorf("msg is not %s", reqMsg)
	}
	return respMsg, nil
}

// ExceptionTest ...
func (m *MockImpl) ExceptionTest(ctx context.Context, req *kt.MockReq) (r string, err error) {
	return "", kt.NewException()
}

func genBinaryResp(method string) []byte {
	idx := 0
	buf := make([]byte, 12+len(method)+len(respMsg))
	binary.BigEndian.PutUint32(buf, thrift.VERSION_1)
	idx += 4
	binary.BigEndian.PutUint32(buf[idx:idx+4], uint32(len(method)))
	idx += 4
	copy(buf[idx:idx+len(method)], method)
	idx += len(method)
	binary.BigEndian.PutUint32(buf[idx:idx+4], 100)
	idx += 4
	copy(buf[idx:idx+len(respMsg)], respMsg)
	return buf
}