go-study/zinx/znet/server_test.go

package znet

import (
	"context"
	"fmt"
	"go-study/zinx/ziface"
	"net"
	"testing"
	"time"
)

type PingRouter struct {
	BaseRouter
}

func (pr *PingRouter) Handle(request ziface.IRequest) {
	fmt.Println("Call PingRouter Handle")
	// 先读取客户端的数据，再回写ping...ping...ping
	fmt.Println("recv from client: msgID=", request.GetMsgID(), ", data=", string(request.GetData()))

	// 回写数据
	err := request.GetConnection().SendMsg(0, []byte("ping...ping...ping\n"))
	if err != nil {
		fmt.Println("call back ping error:", err.Error())
	}
}

type HelloZinxRouter struct {
	BaseRouter
}

func (h *HelloZinxRouter) Handle(request ziface.IRequest) {
	fmt.Println("Call HelloZinxRouter Handle")
	// 先读取客户端的数据，再回写ping...ping...ping
	fmt.Println("recv from client: msgID=", request.GetMsgID(), ", data=", string(request.GetData()))

	// 回写数据
	err := request.GetConnection().SendMsg(1, []byte("Hello Zinx Router V0.6\n"))
	if err != nil {
		fmt.Println("call back HelloZinx error:", err.Error())
	}
}

// DoConnectionBegin 创建连接之后执行的钩子函数
func DoConnectionBegin(conn ziface.IConnection) {
	fmt.Println("DoConnectionBegin is Called ... ")

	ctx := context.WithValue(context.Background(), "Name", "fantasticbin")
	ctx = context.WithValue(ctx, "Home", "https://www.fantasticbin.com")
	conn.SetContext(ctx)
	fmt.Println("Set conn Name, Home done!")

	if err := conn.SendMsg(2, []byte("DoConnection BEGIN")); err != nil {
		fmt.Println("DoConnectionBegin error:", err)
	}
}

// DoConnectionLost 连接断开之前执行的钩子函数
func DoConnectionLost(conn ziface.IConnection) {
	// 断开连接时，获取连接的上下文信息
	ctx := conn.GetContext()
	if ctx != nil {
		if name := ctx.Value("Name"); name != nil {
			fmt.Println("ConnID =", conn.GetConnID(), " get Name from Context =", name)
		}
		if home := ctx.Value("Home"); home != nil {
			fmt.Println("ConnID =", conn.GetConnID(), " get Home from Context =", home)
		}
	}

	fmt.Println("DoConnectionLost is Called ... ")
	fmt.Println("conn ID =", conn.GetConnID(), " is lost...")
}

// ClientTest0 模拟客户端0
func ClientTest0() {
	fmt.Println("Client Test... start")
	// 3s 之后发起测试请求，给服务器端开启服务的机会
	time.Sleep(3 * time.Second)

	conn, err := net.Dial("tcp", "127.0.0.1:7777")
	if err != nil {
		fmt.Println("client dial err:", err)
		return
	}

	for {
		// 发送封包数据
		dp := NewDataPack()
		msg, _ := dp.Pack(NewMsgPackage(0, []byte("Zinx V0.6 Client1 Test Message")))
		if _, err := conn.Write(msg); err != nil {
			fmt.Println("write error:", err)
			return
		}

		// 服务器回复一个数据，先读出流中的 head 部分
		headData := make([]byte, dp.GetHeadLen())
		if _, err := conn.Read(headData); err != nil {
			fmt.Println("read head error:", err)
			break
		}

		// 拆包，得到 msgID 和 dataLen 放在 msg 中
		msgHead, err := dp.Unpack(headData)
		if err != nil {
			fmt.Println("unpack head error:", err)
			return
		}

		if msgHead.GetDataLen() > 0 {
			// msg 是有数据的，需要再次读取 dataLen 个字节的数据
			msg := msgHead.(*Message)
			msg.data = make([]byte, msg.GetDataLen())

			// 根据 dataLen 的长度再次从io中读取
			if _, err := conn.Read(msg.data); err != nil {
				fmt.Println("read msg data error:", err)
				return
			}

			fmt.Println("==> Recv Server Msg: ID=", msg.GetMsgID(), ", len=", msg.GetDataLen(), ", data=", string(msg.GetData()))
		}

		time.Sleep(1 * time.Second)
	}
}

// ClientTest1 模拟客户端1
func ClientTest1() {
	fmt.Println("Client Test... start")
	// 3s 之后发起测试请求，给服务器端开启服务的机会
	time.Sleep(3 * time.Second)

	conn, err := net.Dial("tcp", "127.0.0.1:7777")
	if err != nil {
		fmt.Println("client dial err:", err)
		return
	}

	for {
		// 发送封包数据
		dp := NewDataPack()
		msg, _ := dp.Pack(NewMsgPackage(1, []byte("Zinx V0.6 Client2 Test Message")))
		if _, err := conn.Write(msg); err != nil {
			fmt.Println("write error:", err)
			return
		}

		// 服务器回复一个数据，先读出流中的 head 部分
		headData := make([]byte, dp.GetHeadLen())
		if _, err := conn.Read(headData); err != nil {
			fmt.Println("read head error:", err)
			break
		}

		// 拆包，得到 msgID 和 dataLen 放在 msg 中
		msgHead, err := dp.Unpack(headData)
		if err != nil {
			fmt.Println("unpack head error:", err)
			return
		}

		if msgHead.GetDataLen() > 0 {
			// msg 是有数据的，需要再次读取 dataLen 个字节的数据
			msg := msgHead.(*Message)
			msg.data = make([]byte, msg.GetDataLen())

			// 根据 dataLen 的长度再次从io中读取
			if _, err := conn.Read(msg.data); err != nil {
				fmt.Println("read msg data error:", err)
				return
			}

			fmt.Println("==> Recv Server Msg: ID=", msg.GetMsgID(), ", len=", msg.GetDataLen(), ", data=", string(msg.GetData()))
		}

		time.Sleep(1 * time.Second)
	}
}

// TestServer 服务器端测试
func TestServer(t *testing.T) {
	// 创建一个 Server 句柄
	s := NewServer()

	// 设置连接创建和销毁的钩子函数
	s.SetOnConnStart(DoConnectionBegin)
	s.SetOnConnStop(DoConnectionLost)

	// 给当前 Zinx 框架添加自定义的 Router
	s.AddRouter(0, &PingRouter{})
	s.AddRouter(1, &HelloZinxRouter{})

	// 启动客户端测试0
	go ClientTest0()

	// 启动客户端测试1
	go ClientTest1()

	// 启动服务器
	s.Serve()
}