nnet/test/integration/router_test.go

package integration

import (
	"encoding/json"
	"fmt"
	"net"
	"testing"
	"time"

	"github.com/noahlann/nnet/pkg/nnet"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestRouterFrameData 测试帧数据路由
func TestRouterFrameData(t *testing.T) {
	// 获取随机端口
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	port := listener.Addr().(*net.TCPAddr).Port
	listener.Close()

	cfg := &nnet.Config{
		Addr: fmt.Sprintf("tcp://127.0.0.1:%d", port),
		Codec: &nnet.CodecConfig{
			DefaultCodec:         "json",
			EnableProtocolEncode: false,
		},
	}

	// 创建服务器并在启动前注册路由
	server, err := nnet.NewServer(cfg)
	require.NoError(t, err)

	// 注册帧数据路由
	server.Router().RegisterFrameData("op", "==", "ping", func(ctx nnet.Context) error {
		return ctx.Response().Write(map[string]any{"pong": true})
	})

	server.Router().RegisterFrameData("op", "==", "echo", func(ctx nnet.Context) error {
		return ctx.Response().Write(map[string]any{"echo": string(ctx.Request().Raw())})
	})

	// 启动服务器
	ts := &TestServer{
		Server: server,
		Addr:   cfg.Addr,
		stopCh: make(chan struct{}),
	}

	ts.wg.Add(1)
	go func() {
		defer ts.wg.Done()
		if err := server.Start(); err != nil {
			t.Logf("Server error: %v", err)
		}
	}()

	require.Eventually(t, func() bool {
		return server.Started()
	}, 3*time.Second, 50*time.Millisecond, "Server should start within 3 seconds")
	
	// 等待服务器准备好
	require.Eventually(t, func() bool {
		testConn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 500*time.Millisecond)
		if err != nil {
			return false
		}
		testConn.Close()
		return true
	}, 5*time.Second, 100*time.Millisecond, "Server should be ready")

	defer CleanupTestServer(t, ts)

	client := NewTestClient(t, ts.Addr, nil)
	defer CleanupTestClient(t, client)

	ConnectTestClient(t, client)

	// 等待服务器准备好
	time.Sleep(100 * time.Millisecond)

	// 测试 ping
	pingData := `{"op":"ping"}`
	resp := RequestWithTimeout(t, client, []byte(pingData), 3*time.Second)
	t.Logf("Response for ping: %q", string(resp))
	
	var result map[string]any
	err = json.Unmarshal(resp, &result)
	assert.NoError(t, err, "Response should be valid JSON")
	assert.True(t, result["pong"].(bool), "Response should contain pong: true")

	// 测试 echo
	echoData := `{"op":"echo","msg":"hello"}`
	resp = RequestWithTimeout(t, client, []byte(echoData), 3*time.Second)
	t.Logf("Response for echo: %q", string(resp))
	
	err = json.Unmarshal(resp, &result)
	assert.NoError(t, err, "Response should be valid JSON")
	assert.Contains(t, result["echo"].(string), "echo", "Response should echo the data")
}

// TestRouterMiddleware 测试路由中间件
func TestRouterMiddleware(t *testing.T) {
	// 获取随机端口
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	port := listener.Addr().(*net.TCPAddr).Port
	listener.Close()

	cfg := &nnet.Config{
		Addr: fmt.Sprintf("tcp://127.0.0.1:%d", port),
		Codec: &nnet.CodecConfig{
			DefaultCodec:         "json",
			EnableProtocolEncode: false,
		},
	}

	// 创建服务器并在启动前注册路由
	server, err := nnet.NewServer(cfg)
	require.NoError(t, err)

	// 创建路由分组并添加中间件
	group := server.Router().Group()
	group.Use(func(ctx nnet.Context) error {
		ctx.Set("middleware", "executed")
		return nil
	})

	// 注册路由
	group.RegisterString("hello", func(ctx nnet.Context) error {
		middleware := ctx.GetString("middleware")
		return ctx.Response().Write(map[string]any{
			"message":   "hello with middleware",
			"middleware": middleware,
		})
	})

	// 启动服务器
	ts := &TestServer{
		Server: server,
		Addr:   cfg.Addr,
		stopCh: make(chan struct{}),
	}

	ts.wg.Add(1)
	go func() {
		defer ts.wg.Done()
		if err := server.Start(); err != nil {
			t.Logf("Server error: %v", err)
		}
	}()

	require.Eventually(t, func() bool {
		return server.Started()
	}, 3*time.Second, 50*time.Millisecond, "Server should start within 3 seconds")
	
	// 等待服务器准备好
	require.Eventually(t, func() bool {
		testConn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 500*time.Millisecond)
		if err != nil {
			return false
		}
		testConn.Close()
		return true
	}, 5*time.Second, 100*time.Millisecond, "Server should be ready")

	defer CleanupTestServer(t, ts)

	client := NewTestClient(t, ts.Addr, nil)
	defer CleanupTestClient(t, client)

	ConnectTestClient(t, client)

	// 等待服务器准备好
	time.Sleep(100 * time.Millisecond)

	resp := RequestWithTimeout(t, client, []byte("hello"), 3*time.Second)
	t.Logf("Response for middleware test: %q", string(resp))
	
	var result map[string]any
	err = json.Unmarshal(resp, &result)
	assert.NoError(t, err, "Response should be valid JSON")
	assert.Equal(t, "executed", result["middleware"], "Middleware should be executed")
}