go_dreamfactory/cmd/robot/robot.go

package robot

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go_dreamfactory/comm"
	"go_dreamfactory/pb"
	"io/ioutil"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	jsoniter "github.com/json-iterator/go"
	"google.golang.org/protobuf/proto"
)

type Robot struct {
	ws       *websocket.Conn
	opts     *Options
	user     *pb.DBUser
	builders []*TestCase //测试用例

	linkCase *LinkCase //测试用例链,适应于用例跑批
	wg       sync.WaitGroup
}

func NewRobot(opts *Options) *Robot {
	ws, _, err := websocket.DefaultDialer.Dial(opts.WsUrl, nil)
	if err != nil {
		log.Fatal(err)
	}
	r := &Robot{
		ws:       ws,
		opts:     opts,
		linkCase: NewLinkCase(),
	}

	return r
}

func (r *Robot) Run() {
	log.Print("Robot running...")
	log.Printf("websocket %s \n", r.opts.WsUrl)
	if r.opts.Create { //创建新用户
		r.AccountRegister(r.opts.Account, int32(r.opts.ServerId))
	} else {
		if r.opts.Account == "" {
			log.Fatal("WARNNING: account is required !!!")
		}
		r.AccountLogin()
	}

	//处理响应
	go func() {
		for {			
			var msg *pb.UserMessage = &pb.UserMessage{}
			_, data, err := r.ws.ReadMessage()
			if err != nil {
				log.Println(err)
			}

			if err = proto.Unmarshal(data, msg); err != nil {
				log.Fatal(err)
			}
			r.handleRsp(msg)
		}
	}()

	// select {}
	r.wg.Wait()
}

type TestCase struct {
	Desc      string
	mainType  string
	subType   string
	req       proto.Message
	rsp       proto.Message
	enabled   bool
	start     time.Time
	requested bool //请求标识 true已发
	print     func(rsp proto.Message)
}

func (r *Robot) addBuilders(builders []*TestCase) {
	for _, b := range builders {
		if b.enabled {
			r.builders = append(r.builders, b)
		}
	}
}

//执行请求
func (r *Robot) handleReq() {
	for _, b := range r.builders {
		if b.req != nil && !b.requested {
			r.wg.Add(1)
			time.Sleep(time.Second * 1)
			b.requested = true
			b.start = time.Now()
			head := &pb.UserMessage{MainType: b.mainType, SubType: b.subType}
			defer traceFunc(head.MainType, head.SubType, r.user.GetUid(), b.req)
			err := r.SendToClient(head, b.req)
			if err != nil {
				log.Fatal(err)
			}
		}
	}
}

//执行响应
func (r *Robot) handleRsp(msg *pb.UserMessage) {
	for i, b := range r.builders {
		if b.enabled && (msg.MainType == b.mainType &&
			msg.SubType == b.subType) {
			if !comm.ProtoUnmarshal(msg, b.rsp) {
				return
			}
			if b.print == nil {
				printReply(msg, b)
			} else {
				fmt.Printf("===== %s [%s.%s] =====\n", b.Desc, msg.MainType, msg.SubType)
				b.print(b.rsp)
				fmt.Println("==============================")
			}

			if msg.MainType == "user" && msg.SubType == "login" {
				r.loginCallback(b.rsp)
			} else {
				if b.requested {
					r.builders = append(r.builders[:i], r.builders[i+1:]...)
				}
			}

			r.wg.Done()
		}
	}

}

//登录回调
func (r *Robot) loginCallback(rsp proto.Message) {
	r.builders = append(r.builders[:0], r.builders[1:]...)
	lr := rsp.(*pb.UserLoginResp)
	if lr.Data != nil {
		r.user = lr.Data
		r.onUserLoaded()
	} else {
		r.AccountRegister(r.opts.Account, int32(r.opts.ServerId)) //请求Http接口，模拟创建新账号
	}
}

func (r *Robot) SendToClient(msg *pb.UserMessage, rsp proto.Message) error {
	//模拟客户端 每次请求都会生成新的秘钥
	msg.Sec = r.BuildSecStr()
	if comm.ProtoMarshal(rsp, msg) {
		data, _ := proto.Marshal(msg)
		return r.ws.WriteMessage(websocket.BinaryMessage, data)
	}
	return nil
}

//在这里添加玩家成功登录以后的测试方法
//次方法在用户登录成功后调用
func (r *Robot) onUserLoaded() {
	//notify
	r.RunNotify()
	//user
	r.RunUser()
	//hero
	r.RunHero()
	//friend
	r.RunFriend()

	//pack
	// r.RunPack()

}

//注册账号
func (r *Robot) AccountRegister(account string, sid int32) {
	if account == "" {
		log.Fatal("account value is empty")
	}
	//http
	regReq := &pb.UserRegisterReq{Account: account, Sid: sid}
	jsonByte, _ := json.Marshal(regReq)
	req, err := http.NewRequest("POST", r.opts.RegUrl, bytes.NewReader(jsonByte))
	if err != nil {
		log.Fatalf("account register err %v", err)
	}
	req.Header.Set("Content-Type", "application/json;charset=UTF-8")
	httpClient := &http.Client{}
	rsp, err := httpClient.Do(req)
	if err != nil {
		panic(err)
	}
	defer rsp.Body.Close()

	body, _ := ioutil.ReadAll(rsp.Body)
	regRsp := &pb.UserRegisterResp{}
	err = jsoniter.Unmarshal(body, regRsp)

	if regRsp.Code == pb.ErrorCode_Success { //注册成功

		fmt.Printf("account:%s 注册成功", regRsp.Account)
		//登录
		var user_builders = []*TestCase{
			{
				Desc:     "登录",
				mainType: "user",
				subType:  "login",
				req: &pb.UserLoginReq{
					Account: account,
					Sid:     sid,
				},
				rsp:     &pb.UserLoginResp{},
				enabled: true,
			},
		}
		r.addBuilders(user_builders)
		r.handleReq()
	}
}

//打印响应
func printReply(msg *pb.UserMessage, builder *TestCase) {
	if m, ok := builder.rsp.(*pb.NotifyErrorNotifyPush); ok {
		var tt time.Duration
		if builder.start.IsZero() {
			tt = time.Duration(0)
		} else {
			tt = time.Since(builder.start)
		}
		log.Printf("rsp %s [%v] [%s.%s] [%v:%v]", builder.Desc, tt, m.ReqMainType, m.ReqSubType, int32(m.Code), m.Data)
	} else {
		log.Printf("rsp %s [%v] [%s.%s] [%v]", builder.Desc, time.Since(builder.start), msg.MainType, msg.SubType, builder.rsp)
	}
}

//方法参数跟踪
func traceFunc(module string, funcName string, uid string, funcArgs interface{}) {
	log.Printf("req [%s.%s] [%s] [%v]", module, funcName, uid, funcArgs)
}