zogo/spmux/spmux_test.go

package spmux

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"net/http"
	"testing"
	"time"
)

// freePort 找一个空闲端口（有轻微竞态，测试环境可接受）
func freePort(t *testing.T) int {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	port := l.Addr().(*net.TCPAddr).Port
	l.Close()
	return port
}

func TestPortMuxHttpRouting(t *testing.T) {
	port := freePort(t)

	pm := NewPortMux(port, nil)
	if err := pm.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}

	httpDone := make(chan struct{})
	go func() {
		defer close(httpDone)
		http.Serve(pm.HttpListener(), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			fmt.Fprint(w, "mux-http-ok")
		}))
	}()
	// 先 Close 解除 Serve 阻塞，再等待协程退出（顺序不能反，否则死锁）
	defer func() {
		pm.Close()
		<-httpDone
	}()

	time.Sleep(100 * time.Millisecond)

	resp, err := http.Get(fmt.Sprintf("http://127.0.0.1:%d/", port))
	if err != nil {
		t.Fatalf("HTTP 请求失败: %v", err)
	}
	body, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if string(body) != "mux-http-ok" {
		t.Errorf("HTTP 响应 = %q", string(body))
	}
}

func TestPortMuxTcpRouting(t *testing.T) {
	port := freePort(t)

	pm := NewPortMux(port, nil)
	if err := pm.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}

	tcpDone := make(chan struct{})
	go func() {
		defer close(tcpDone)
		for {
			c, err := pm.TcpListener().Accept()
			if err != nil {
				return
			}
			go func(cc net.Conn) {
				defer cc.Close()
				buf := make([]byte, 1024)
				for {
					n, err := cc.Read(buf)
					if n > 0 {
						cc.Write(append(buf[:n:n], '\n'))
					}
					if err != nil {
						return
					}
				}
			}(c)
		}
	}()
	// 先 Close 解除 Accept 阻塞，再等待协程退出（顺序不能反，否则死锁）
	defer func() {
		pm.Close()
		<-tcpDone
	}()

	time.Sleep(100 * time.Millisecond)

	// 非 HTTP 数据（超过 8 字节）应路由到 TcpListener
	raw, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", port))
	if err != nil {
		t.Fatalf("TCP 连接失败: %v", err)
	}
	fmt.Fprint(raw, "ping-ping")
	raw.SetReadDeadline(time.Now().Add(3 * time.Second))
	line, err := bufio.NewReader(raw).ReadString('\n')
	if err != nil {
		t.Fatalf("读取 echo: %v", err)
	}
	if line != "ping-ping\n" {
		t.Errorf("echo = %q", line)
	}
	raw.Close()
}

func TestPortListenerAddr(t *testing.T) {
	port := freePort(t)
	pm := NewPortMux(port, nil)
	if err := pm.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer pm.Close()

	if pm.HttpListener().Addr().String() == "" {
		t.Error("HttpListener Addr 不应为空")
	}
	if pm.TcpListener().Addr().String() == "" {
		t.Error("TcpListener Addr 不应为空")
	}
}