go_dreamfactory/stress/robot/robot.go

package robot

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go_dreamfactory/comm"
	"io"
	"os"

	"go_dreamfactory/pb"

	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/proto"
)

type Robot struct {
	ws      *websocket.Conn
	account string
	sid     string
	user    *pb.DBUser
}

var zlog *logrus.Logger

func initlog() {
	zlog = logrus.New()
	zlog.SetLevel(logrus.DebugLevel)
	zlog.SetFormatter(&RobotFormatter{
		DisableTimestamp: true,
		DisableQuote:     true,
	})

	//设置output,默认为stderr,可以为任何io.Writer，比如文件*os.File
	file, err := os.OpenFile("robot.log", os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
	writers := []io.Writer{
		file,
		os.Stdout}

	//同时写文件和屏幕
	fileAndStdoutWriter := io.MultiWriter(writers...)
	if err == nil {
		zlog.SetOutput(fileAndStdoutWriter)
	} else {
		zlog.Info("failed to log to file.")
	}
}

func NewRobot(url string) *Robot {
	initlog()
	ws, _, err := websocket.DefaultDialer.Dial(url, nil)
	if err != nil {
		zlog.Fatalf("websocket conn err:%v", err)
	}
	r := &Robot{
		ws:      ws,
		account: "user001",
		sid:     "dfmxf",
	}

	return r
}

type TestCase struct {
	id       string        //用例ID 如果没有指定,会自动赋值uuid
	desc     string        //用例描述
	mainType string        //协议类型 L1
	subType  string        //协议类型 L2
	req      proto.Message //请求类型
	rsp      proto.Message //响应类型
}

//加入用例并执行请求
func (r *Robot) addTestCaseAndReq(tcs []*TestCase) {

}

//错误通知处理
func (r *Robot) handleNotify(uuid string, msg *pb.UserMessage) {
	if msg.MainType == "notify" && msg.SubType == "errornotify" {

		rsp := &pb.NotifyErrorNotifyPush{}
		if !comm.ProtoUnmarshal(msg, rsp) {
			return
		}
	}
}

//处理响应
func (r *Robot) handleRsp(id string) {

	var msg *pb.UserMessage = &pb.UserMessage{}
	_, data, err := r.ws.ReadMessage()
	if err != nil {
		zlog.Fatalf("readMessage err:%v", err)
	}

	if err = proto.Unmarshal(data, msg); err != nil {
		zlog.Fatalf("unmarshal err:%v", err)
	}

}

//登录回调
func (r *Robot) loginCallback(rsp proto.Message) {
	//清除登录用例
	lr := rsp.(*pb.UserLoginResp)
	if lr.Data != nil {
		r.user = lr.Data
		r.onUserLoaded()
	}
}

//发送消息
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
}

//格式化json
func formatJson(data string) (string, error) {
	var out bytes.Buffer
	err := json.Indent(&out, []byte(data), "", "    ")
	return out.String(), err
}

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

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

type RobotFormatter struct {
	DisableTimestamp bool
	DisableQuote     bool
}

func (r *RobotFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	var timestamp string
	var newLog string
	if !r.DisableTimestamp {
		timestamp = entry.Time.Format("2006-01-02 15:04:05")
		newLog = fmt.Sprintf("[%s] %s\n", timestamp, entry.Message)
	} else {
		newLog = fmt.Sprintf("%s\n", entry.Message)
	}

	b.WriteString(newLog)
	return b.Bytes(), nil
}

func (r *Robot) MessageRsp(mainType, subType string) (bool, int64) {
	var byteLen int64
	var msg *pb.UserMessage = &pb.UserMessage{}
	_, data, err := r.ws.ReadMessage()
	if err != nil {
		fmt.Printf("readMessage err:%v", err)
		return false, byteLen
	}
	byteLen = int64(len(data))
	if err = proto.Unmarshal(data, msg); err != nil {
		fmt.Printf("unmarshal err:%v", err)
	}
	//fmt.Printf("接收消息=====msg.MainType = %s,msg.SubType = %s\n", msg.MainType, msg.SubType)
	if msg.MainType == "chat" && msg.SubType == "message" {
		resp := &pb.ChatMessagePush{}
		if !comm.ProtoUnmarshal(msg, resp) {
			return false, byteLen
		}
		return false, byteLen
		//fmt.Printf("接收消息=====resp:%v\n", resp.Chat)
	} else if msg.MainType == mainType && msg.SubType == subType {
		resp := &pb.UserLoginResp{}
		if !comm.ProtoUnmarshal(msg, resp) {
			return false, byteLen
		}
		//fmt.Printf("接收消息=====loginData:%v,userExpand:%v\n", resp.Data, resp.Ex)
		return true, byteLen
	}
	return false, byteLen
}

func (r *Robot) CloseHandler() {

	r.ws.Close()

}

func (r *Robot) GetMessagedata() []byte {
	_, data, err := r.ws.ReadMessage()
	if err != nil {
		fmt.Printf("readMessage err:%v", err)
	}
	return data
}
func (r *Robot) Login() error {
	msg := &pb.UserMessage{MainType: "user", SubType: "login"}
	rsp := &pb.UserLoginReq{
		Account: "user001",
		Sid:     "dfmxf",
	}
	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) Create() bool {
	mainType := "user"
	subType := "create"
	msg := &pb.UserMessage{MainType: mainType, SubType: subType}
	rsp := &pb.UserCreateReq{
		NickName: "001",
		Figure:   100,
		Gender:   0,
	}
	msg.Sec = r.BuildSecStr()
	if comm.ProtoMarshal(rsp, msg) {
		data, _ := proto.Marshal(msg)
		err := r.ws.WriteMessage(websocket.BinaryMessage, data)
		if err != nil {
			fmt.Printf("WriteMessage err:%v", err)
		}
	}

	for {
		if b, _ := r.MessageRsp(mainType, subType); b {
			return true
		}
	}

}