github.com/cloudwego/kitex@v0.9.0/pkg/generic/reflect_test/map_test.go

/**
 * Copyright 2023 ByteDance Inc.
 *
 * 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

import (
	"context"
	"errors"
	"math/rand"
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/cloudwego/kitex/client"
	"github.com/cloudwego/kitex/client/callopt"
	"github.com/cloudwego/kitex/client/genericclient"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/generic"
	"github.com/cloudwego/kitex/server"
	"github.com/cloudwego/kitex/server/genericserver"
)

var mcli genericclient.Client

func TestThriftMapExample(t *testing.T) {
	testThriftMapExample(t)
}

func BenchmarkThriftMapExample(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testThriftMapExample(b)
	}
}

func testThriftMapExample(t testing.TB) {
	log_id := strconv.Itoa(rand.Int())

	req := makeExampleReqMap(true, reqMsg, log_id)

	out, err := mcli.GenericCall(context.Background(), method, req, callopt.WithRPCTimeout(100*time.Second))
	test.Assert(t, err == nil, err)
	resp, ok := out.(map[string]interface{})
	test.Assert(t, ok)

	msg, ok := resp["Msg"].(string)
	if !ok || msg != respMsg {
		t.Fail()
	}
	require_field, ok := resp["required_field"].(string)
	if !ok || require_field != reqMsg {
		t.Fail()
	}
	logid, ok := resp["BaseResp"].(map[string]interface{})["StatusMessage"]
	if !ok || log_id != logid {
		t.Fail()
	}
}

func initThriftMapClient(addr, idl string) genericclient.Client {
	p, err := generic.NewThriftFileProvider(idl)
	if err != nil {
		panic(err)
	}
	g, err := generic.MapThriftGeneric(p)
	if err != nil {
		panic(err)
	}
	cli := newGenericClient("destServiceName", g, addr)
	if err != nil {
		panic(err)
	}
	return cli
}

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

func initThriftMapServer(address, idl string, handler generic.Service) server.Server {
	addr, _ := net.ResolveTCPAddr("tcp", address)
	p, err := generic.NewThriftFileProvider(idl)
	if err != nil {
		panic(err)
	}
	g, err := generic.MapThriftGeneric(p)
	if err != nil {
		panic(err)
	}
	svr := newGenericServer(g, addr, handler)
	return svr
}

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
}

type exampeServerImpl struct{}

// GenericCall ...
func (g *exampeServerImpl) GenericCall(ctx context.Context, method string, request interface{}) (response interface{}, err error) {
	buf := request.(map[string]interface{})

	required_field := ""
	logid := ""
	if b, ok := buf["B"].(bool); ok && b {
		if base, ok := buf["Base"].(map[string]interface{}); !ok {
			return nil, errors.New("empty base field")
		} else {
			logid, _ = base["LogID"].(string)
		}
		if testmap, ok := buf["TestMap"].(map[interface{}]interface{}); !ok {
			return nil, errors.New("empty TestMap")
		} else if a, ok := testmap["a"].(map[string]interface{}); !ok {
			return nil, errors.New("key 'a' not found")
		} else {
			required_field, _ = a["Bar"].(string)
		}
	}

	return makeExampleRespMap(respMsg, required_field, logid)
}

func makeExampleRespMap(msg, require_field, logid string) (map[string]interface{}, error) {
	return map[string]interface{}{
		"Msg":            msg,
		"required_field": require_field,
		"BaseResp": map[string]interface{}{
			"StatusMessage": logid,
		},
	}, nil
}

func makeExampleReqMap(B bool, A, logid string) map[string]interface{} {
	list := make([]interface{}, SampleListSize+1)
	list[0] = map[string]interface{}{
		"Bar": A,
	}
	for i := 1; i < len(list); i++ {
		list[i] = map[string]interface{}{
			"Bar": A,
		}
	}
	m := make(map[string]interface{}, SampleListSize+1)
	m["a"] = map[string]interface{}{
		"Bar": A,
	}
	for i := 1; i < len(list); i++ {
		m[strconv.Itoa(i)] = map[string]interface{}{
			"Bar": A,
		}
	}
	return map[string]interface{}{
		"Msg":      "Hello",
		"Foo":      int32(1),
		"TestList": list,
		"TestMap":  m,
		"I64List": []interface{}{
			int64(1), int64(2), int64(3),
		},
		"B": B,
		"Base": map[string]interface{}{
			"LogID":  logid,
			"Caller": "d.e.f",
			"Addr":   "127.0.0.1",
			"Client": "dynamicgo",
		},
	}
}