xfy
6ae7e32ef1
实现 Phase 3 核心功能: - loadbalance: 轮询、加权轮询、最少连接、IP哈希四种算法 - proxy: HTTP 反向代理、健康检查、故障转移 - 所有实现均为并发安全，使用 atomic 操作 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
899 lines
22 KiB
Go
899 lines
22 KiB
Go
// Package proxy provides reverse proxy functionality for the Lolly HTTP server.
|
||
//
|
||
//go:generate go test -v ./...
|
||
package proxy
|
||
|
||
import (
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/valyala/fasthttp"
|
||
"github.com/valyala/fasthttp/fasthttputil"
|
||
|
||
"rua.plus/lolly/internal/config"
|
||
"rua.plus/lolly/internal/loadbalance"
|
||
)
|
||
|
||
// TestNewProxy 测试 NewProxy 函数
|
||
func TestNewProxy(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
cfg *config.ProxyConfig
|
||
targets []*loadbalance.Target
|
||
wantErr bool
|
||
errContains string
|
||
}{
|
||
{
|
||
name: "正常创建",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second, Read: 30 * time.Second, Write: 30 * time.Second},
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: true},
|
||
{URL: "http://localhost:8082", Healthy: true},
|
||
},
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "nil配置",
|
||
cfg: nil,
|
||
targets: []*loadbalance.Target{{URL: "http://localhost:8081"}},
|
||
wantErr: true,
|
||
errContains: "proxy config is nil",
|
||
},
|
||
{
|
||
name: "空目标列表",
|
||
cfg: &config.ProxyConfig{Path: "/api"},
|
||
targets: []*loadbalance.Target{},
|
||
wantErr: true,
|
||
errContains: "no proxy targets provided",
|
||
},
|
||
{
|
||
name: "nil目标列表",
|
||
cfg: &config.ProxyConfig{Path: "/api"},
|
||
targets: nil,
|
||
wantErr: true,
|
||
errContains: "no proxy targets provided",
|
||
},
|
||
{
|
||
name: "默认负载均衡算法",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "",
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: true},
|
||
},
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "加权轮询算法",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "weighted_round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Weight: 1, Healthy: true},
|
||
{URL: "http://localhost:8082", Weight: 2, Healthy: true},
|
||
},
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "最少连接算法",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "least_conn",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: true},
|
||
},
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "IP哈希算法",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "ip_hash",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: true},
|
||
},
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "无效负载均衡算法",
|
||
cfg: &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "invalid_algorithm",
|
||
},
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: true},
|
||
},
|
||
wantErr: true,
|
||
errContains: "unsupported load balance algorithm",
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
p, err := NewProxy(tt.cfg, tt.targets)
|
||
if tt.wantErr {
|
||
if err == nil {
|
||
t.Errorf("NewProxy() expected error containing %q, got nil", tt.errContains)
|
||
return
|
||
}
|
||
if !contains(err.Error(), tt.errContains) {
|
||
t.Errorf("NewProxy() error = %v, want containing %q", err, tt.errContains)
|
||
}
|
||
return
|
||
}
|
||
if err != nil {
|
||
t.Errorf("NewProxy() unexpected error: %v", err)
|
||
return
|
||
}
|
||
if p == nil {
|
||
t.Error("NewProxy() returned nil proxy")
|
||
return
|
||
}
|
||
if p.config != tt.cfg {
|
||
t.Error("NewProxy() proxy config not set correctly")
|
||
}
|
||
if p.balancer == nil {
|
||
t.Error("NewProxy() balancer not initialized")
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestServeHTTP_NoHealthyTargets 测试没有健康目标时返回502
|
||
func TestServeHTTP_NoHealthyTargets(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second, Read: 30 * time.Second, Write: 30 * time.Second},
|
||
}
|
||
|
||
// 所有目标都不健康
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8081", Healthy: false},
|
||
{URL: "http://localhost:8082", Healthy: false},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
// 创建测试请求
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.Header.SetMethod(fasthttp.MethodGet)
|
||
ctx.Request.SetRequestURI("/api/test")
|
||
|
||
// 执行请求
|
||
p.ServeHTTP(ctx)
|
||
|
||
// 应该返回502
|
||
if ctx.Response.StatusCode() != fasthttp.StatusBadGateway {
|
||
t.Errorf("ServeHTTP() status code = %d, want %d", ctx.Response.StatusCode(), fasthttp.StatusBadGateway)
|
||
}
|
||
}
|
||
|
||
// TestServeHTTP_RequestForwarding 测试请求转发
|
||
func TestServeHTTP_RequestForwarding(t *testing.T) {
|
||
// 创建本地测试服务器
|
||
ln := fasthttputil.NewInmemoryListener()
|
||
defer ln.Close()
|
||
|
||
// 启动后端服务器
|
||
go func() {
|
||
s := &fasthttp.Server{
|
||
Handler: func(ctx *fasthttp.RequestCtx) {
|
||
ctx.SetStatusCode(fasthttp.StatusOK)
|
||
ctx.SetBodyString("Hello from backend")
|
||
ctx.Response.Header.Set("X-Backend-Header", "test-value")
|
||
},
|
||
}
|
||
s.Serve(ln)
|
||
}()
|
||
|
||
// 等待服务器启动
|
||
time.Sleep(10 * time.Millisecond)
|
||
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second, Read: 30 * time.Second, Write: 30 * time.Second},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
// 创建测试请求
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.Header.SetMethod(fasthttp.MethodGet)
|
||
ctx.Request.SetRequestURI("/api/test")
|
||
ctx.Request.Header.Set("X-Custom-Header", "client-value")
|
||
|
||
// 执行请求
|
||
p.ServeHTTP(ctx)
|
||
|
||
// 由于没有真实后端，应该返回502
|
||
// 但在单元测试中我们可以验证错误处理逻辑
|
||
if ctx.Response.StatusCode() != fasthttp.StatusBadGateway {
|
||
t.Logf("ServeHTTP() status code = %d (expected 502 when no backend available)", ctx.Response.StatusCode())
|
||
}
|
||
}
|
||
|
||
// TestSelectTarget 测试目标选择
|
||
func TestSelectTarget(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
loadBalance string
|
||
targets []*loadbalance.Target
|
||
clientIP string
|
||
expectedTarget string
|
||
}{
|
||
{
|
||
name: "轮询选择",
|
||
loadBalance: "round_robin",
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://backend1:8080", Healthy: true},
|
||
{URL: "http://backend2:8080", Healthy: true},
|
||
},
|
||
expectedTarget: "http://backend1:8080",
|
||
},
|
||
{
|
||
name: "跳过不健康目标",
|
||
loadBalance: "round_robin",
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://backend1:8080", Healthy: false},
|
||
{URL: "http://backend2:8080", Healthy: true},
|
||
},
|
||
expectedTarget: "http://backend2:8080",
|
||
},
|
||
{
|
||
name: "IP哈希选择",
|
||
loadBalance: "ip_hash",
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://backend1:8080", Healthy: true},
|
||
{URL: "http://backend2:8080", Healthy: true},
|
||
},
|
||
clientIP: "192.168.1.100",
|
||
expectedTarget: "any", // IP哈希应该返回一个目标，具体是哪个取决于哈希值
|
||
},
|
||
{
|
||
name: "所有目标都不健康",
|
||
loadBalance: "round_robin",
|
||
targets: []*loadbalance.Target{
|
||
{URL: "http://backend1:8080", Healthy: false},
|
||
{URL: "http://backend2:8080", Healthy: false},
|
||
},
|
||
expectedTarget: "",
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: tt.loadBalance,
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, tt.targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
if tt.clientIP != "" {
|
||
// 设置远程地址模拟客户端IP
|
||
ctx.Request.Header.Set("X-Forwarded-For", tt.clientIP)
|
||
}
|
||
ctx.Request.SetRequestURI("/api/test")
|
||
|
||
target := p.selectTarget(ctx)
|
||
|
||
if tt.expectedTarget == "" {
|
||
if target != nil {
|
||
t.Errorf("selectTarget() expected nil, got %v", target.URL)
|
||
}
|
||
return
|
||
}
|
||
|
||
if tt.loadBalance == "round_robin" && tt.expectedTarget != "" {
|
||
// 轮询应该选择第一个健康目标
|
||
if target == nil {
|
||
t.Error("selectTarget() returned nil for healthy targets")
|
||
return
|
||
}
|
||
if target.URL != tt.expectedTarget {
|
||
t.Errorf("selectTarget() = %v, want %v", target.URL, tt.expectedTarget)
|
||
}
|
||
}
|
||
|
||
// IP哈希应该始终返回同一个目标给同一个IP
|
||
if tt.loadBalance == "ip_hash" && tt.clientIP != "" {
|
||
if target == nil {
|
||
t.Error("selectTarget() returned nil for IP hash")
|
||
return
|
||
}
|
||
// 再次选择，应该返回相同的目标
|
||
target2 := p.selectTarget(ctx)
|
||
if target2 == nil || target2.URL != target.URL {
|
||
t.Error("IP hash should consistently return the same target for the same IP")
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestModifyRequestHeaders 测试请求头修改
|
||
func TestModifyRequestHeaders(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
clientIP string
|
||
existingXFF string
|
||
setRequest map[string]string
|
||
removeHeaders []string
|
||
checkHeaders map[string]string
|
||
shouldNotExist []string
|
||
}{
|
||
{
|
||
name: "设置X-Real-IP",
|
||
clientIP: "192.168.1.100",
|
||
checkHeaders: map[string]string{
|
||
"X-Real-IP": "192.168.1.100",
|
||
},
|
||
},
|
||
{
|
||
name: "追加X-Forwarded-For",
|
||
clientIP: "192.168.1.100",
|
||
existingXFF: "10.0.0.1",
|
||
checkHeaders: map[string]string{
|
||
"X-Forwarded-For": "10.0.0.1, 10.0.0.1",
|
||
},
|
||
},
|
||
{
|
||
name: "新建X-Forwarded-For",
|
||
clientIP: "192.168.1.100",
|
||
checkHeaders: map[string]string{
|
||
"X-Forwarded-For": "192.168.1.100",
|
||
},
|
||
},
|
||
{
|
||
name: "自定义请求头",
|
||
clientIP: "192.168.1.100",
|
||
setRequest: map[string]string{
|
||
"X-Custom-Header": "custom-value",
|
||
"X-Another": "another-value",
|
||
},
|
||
checkHeaders: map[string]string{
|
||
"X-Custom-Header": "custom-value",
|
||
"X-Another": "another-value",
|
||
},
|
||
},
|
||
{
|
||
name: "移除请求头",
|
||
clientIP: "192.168.1.100",
|
||
removeHeaders: []string{"X-Remove-Me"},
|
||
shouldNotExist: []string{"X-Remove-Me"},
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
Headers: config.ProxyHeaders{
|
||
SetRequest: tt.setRequest,
|
||
Remove: tt.removeHeaders,
|
||
},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.SetRequestURI("/api/test")
|
||
|
||
// 设置客户端IP
|
||
if tt.clientIP != "" {
|
||
ctx.Request.Header.Set("X-Real-IP", tt.clientIP)
|
||
}
|
||
|
||
// 设置已有的X-Forwarded-For
|
||
if tt.existingXFF != "" {
|
||
ctx.Request.Header.Set("X-Forwarded-For", tt.existingXFF)
|
||
}
|
||
|
||
// 设置需要被移除的头
|
||
if len(tt.removeHeaders) > 0 {
|
||
for _, h := range tt.removeHeaders {
|
||
ctx.Request.Header.Set(h, "should-be-removed")
|
||
}
|
||
}
|
||
|
||
target := &loadbalance.Target{URL: "http://localhost:8080"}
|
||
p.modifyRequestHeaders(ctx, target)
|
||
|
||
// 检查期望存在的头
|
||
for key, expectedValue := range tt.checkHeaders {
|
||
actualValue := string(ctx.Request.Header.Peek(key))
|
||
if actualValue != expectedValue {
|
||
t.Errorf("Header %s = %q, want %q", key, actualValue, expectedValue)
|
||
}
|
||
}
|
||
|
||
// 检查不应该存在的头
|
||
for _, key := range tt.shouldNotExist {
|
||
if ctx.Request.Header.Peek(key) != nil {
|
||
t.Errorf("Header %s should not exist", key)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestModifyResponseHeaders 测试响应头修改
|
||
func TestModifyResponseHeaders(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
setResponse map[string]string
|
||
checkHeaders map[string]string
|
||
}{
|
||
{
|
||
name: "设置自定义响应头",
|
||
setResponse: map[string]string{
|
||
"X-Custom-Response": "custom-value",
|
||
"X-Powered-By": "Lolly",
|
||
},
|
||
checkHeaders: map[string]string{
|
||
"X-Custom-Response": "custom-value",
|
||
"X-Powered-By": "Lolly",
|
||
},
|
||
},
|
||
{
|
||
name: "空响应头配置",
|
||
setResponse: nil,
|
||
checkHeaders: map[string]string{},
|
||
},
|
||
{
|
||
name: "覆盖已有响应头",
|
||
setResponse: map[string]string{
|
||
"Content-Type": "application/json",
|
||
},
|
||
checkHeaders: map[string]string{
|
||
"Content-Type": "application/json",
|
||
},
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
Headers: config.ProxyHeaders{
|
||
SetResponse: tt.setResponse,
|
||
},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Response.SetStatusCode(fasthttp.StatusOK)
|
||
|
||
p.modifyResponseHeaders(ctx)
|
||
|
||
// 检查期望存在的头
|
||
for key, expectedValue := range tt.checkHeaders {
|
||
actualValue := string(ctx.Response.Header.Peek(key))
|
||
if actualValue != expectedValue {
|
||
t.Errorf("Response Header %s = %q, want %q", key, actualValue, expectedValue)
|
||
}
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestGetClientIP 测试客户端IP提取
|
||
func TestGetClientIP(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
xff string
|
||
xri string
|
||
expected string
|
||
}{
|
||
{
|
||
name: "从X-Forwarded-For提取",
|
||
xff: "10.0.0.1, 10.0.0.2",
|
||
expected: "10.0.0.1",
|
||
},
|
||
{
|
||
name: "从X-Real-IP提取",
|
||
xri: "192.168.1.100",
|
||
expected: "192.168.1.100",
|
||
},
|
||
{
|
||
name: "X-Forwarded-For优先",
|
||
xff: "10.0.0.1",
|
||
xri: "192.168.1.100",
|
||
expected: "10.0.0.1",
|
||
},
|
||
{
|
||
name: "单IP",
|
||
xff: "10.0.0.1",
|
||
expected: "10.0.0.1",
|
||
},
|
||
{
|
||
name: "带空格",
|
||
xff: " 10.0.0.1 ",
|
||
expected: "10.0.0.1",
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
ctx := &fasthttp.RequestCtx{}
|
||
if tt.xff != "" {
|
||
ctx.Request.Header.Set("X-Forwarded-For", tt.xff)
|
||
}
|
||
if tt.xri != "" {
|
||
ctx.Request.Header.Set("X-Real-IP", tt.xri)
|
||
}
|
||
|
||
ip := getClientIP(ctx)
|
||
if ip != tt.expected {
|
||
t.Errorf("getClientIP() = %q, want %q", ip, tt.expected)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestUpdateTargets 测试更新目标
|
||
func TestUpdateTargets(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
}
|
||
|
||
initialTargets := []*loadbalance.Target{
|
||
{URL: "http://old1:8080", Healthy: true},
|
||
{URL: "http://old2:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, initialTargets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
// 更新目标
|
||
newTargets := []*loadbalance.Target{
|
||
{URL: "http://new1:8080", Healthy: true},
|
||
{URL: "http://new2:8080", Healthy: true},
|
||
{URL: "http://new3:8080", Healthy: true},
|
||
}
|
||
|
||
err = p.UpdateTargets(newTargets)
|
||
if err != nil {
|
||
t.Errorf("UpdateTargets() error: %v", err)
|
||
}
|
||
|
||
// 验证目标已更新
|
||
targets := p.GetTargets()
|
||
if len(targets) != len(newTargets) {
|
||
t.Errorf("UpdateTargets() targets count = %d, want %d", len(targets), len(newTargets))
|
||
}
|
||
|
||
// 验证空目标列表返回错误
|
||
err = p.UpdateTargets([]*loadbalance.Target{})
|
||
if err == nil {
|
||
t.Error("UpdateTargets([]) should return error")
|
||
}
|
||
|
||
// 验证nil目标列表返回错误
|
||
err = p.UpdateTargets(nil)
|
||
if err == nil {
|
||
t.Error("UpdateTargets(nil) should return error")
|
||
}
|
||
}
|
||
|
||
// TestGetTargets 测试获取目标列表
|
||
func TestGetTargets(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://backend1:8080", Healthy: true},
|
||
{URL: "http://backend2:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
gotTargets := p.GetTargets()
|
||
if len(gotTargets) != len(targets) {
|
||
t.Errorf("GetTargets() returned %d targets, want %d", len(gotTargets), len(targets))
|
||
}
|
||
|
||
for i, target := range gotTargets {
|
||
if target.URL != targets[i].URL {
|
||
t.Errorf("GetTargets()[%d].URL = %q, want %q", i, target.URL, targets[i].URL)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestGetConfig 测试获取配置
|
||
func TestGetConfig(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
gotConfig := p.GetConfig()
|
||
if gotConfig != cfg {
|
||
t.Error("GetConfig() returned different config")
|
||
}
|
||
|
||
if gotConfig.Path != cfg.Path {
|
||
t.Errorf("GetConfig().Path = %q, want %q", gotConfig.Path, cfg.Path)
|
||
}
|
||
}
|
||
|
||
// TestIsWebSocketRequest 测试WebSocket请求检测
|
||
func TestIsWebSocketRequest(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
upgrade string
|
||
connection string
|
||
expected bool
|
||
}{
|
||
{
|
||
name: "标准WebSocket请求",
|
||
upgrade: "websocket",
|
||
connection: "upgrade",
|
||
expected: true,
|
||
},
|
||
{
|
||
name: "大小写不敏感",
|
||
upgrade: "WebSocket",
|
||
connection: "Upgrade",
|
||
expected: true,
|
||
},
|
||
{
|
||
name: "非WebSocket升级",
|
||
upgrade: "h2c",
|
||
connection: "upgrade",
|
||
expected: false,
|
||
},
|
||
{
|
||
name: "非upgrade连接",
|
||
upgrade: "websocket",
|
||
connection: "keep-alive",
|
||
expected: false,
|
||
},
|
||
{
|
||
name: "keep-alive, Upgrade",
|
||
upgrade: "websocket",
|
||
connection: "keep-alive, Upgrade",
|
||
expected: true,
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
ctx := &fasthttp.RequestCtx{}
|
||
if tt.upgrade != "" {
|
||
ctx.Request.Header.Set("Upgrade", tt.upgrade)
|
||
}
|
||
if tt.connection != "" {
|
||
ctx.Request.Header.Set("Connection", tt.connection)
|
||
}
|
||
|
||
result := isWebSocketRequest(ctx)
|
||
if result != tt.expected {
|
||
t.Errorf("isWebSocketRequest() = %v, want %v", result, tt.expected)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestCreateBalancer 测试负载均衡器创建
|
||
func TestCreateBalancer(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
algorithm string
|
||
wantErr bool
|
||
errContains string
|
||
}{
|
||
{
|
||
name: "轮询",
|
||
algorithm: "round_robin",
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "加权轮询",
|
||
algorithm: "weighted_round_robin",
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "最少连接",
|
||
algorithm: "least_conn",
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "IP哈希",
|
||
algorithm: "ip_hash",
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "空算法（默认轮询）",
|
||
algorithm: "",
|
||
wantErr: false,
|
||
},
|
||
{
|
||
name: "无效算法",
|
||
algorithm: "unknown_algorithm",
|
||
wantErr: true,
|
||
errContains: "unsupported load balance algorithm",
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
balancer, err := createBalancer(tt.algorithm)
|
||
if tt.wantErr {
|
||
if err == nil {
|
||
t.Errorf("createBalancer(%q) expected error", tt.algorithm)
|
||
return
|
||
}
|
||
if !contains(err.Error(), tt.errContains) {
|
||
t.Errorf("createBalancer(%q) error = %v, want containing %q", tt.algorithm, err, tt.errContains)
|
||
}
|
||
return
|
||
}
|
||
if err != nil {
|
||
t.Errorf("createBalancer(%q) unexpected error: %v", tt.algorithm, err)
|
||
return
|
||
}
|
||
if balancer == nil {
|
||
t.Errorf("createBalancer(%q) returned nil balancer", tt.algorithm)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestCreateHostClient 测试HostClient创建
|
||
func TestCreateHostClient(t *testing.T) {
|
||
tests := []struct {
|
||
name string
|
||
targetURL string
|
||
timeout config.ProxyTimeout
|
||
}{
|
||
{
|
||
name: "HTTP地址",
|
||
targetURL: "http://localhost:8080",
|
||
timeout: config.ProxyTimeout{Connect: 5 * time.Second, Read: 30 * time.Second, Write: 30 * time.Second},
|
||
},
|
||
{
|
||
name: "HTTPS地址",
|
||
targetURL: "https://localhost:8443",
|
||
timeout: config.ProxyTimeout{Connect: 5 * time.Second, Read: 30 * time.Second, Write: 30 * time.Second},
|
||
},
|
||
{
|
||
name: "带路径的URL",
|
||
targetURL: "http://localhost:8080/path",
|
||
timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
},
|
||
}
|
||
|
||
for _, tt := range tests {
|
||
t.Run(tt.name, func(t *testing.T) {
|
||
client := createHostClient(tt.targetURL, tt.timeout)
|
||
if client == nil {
|
||
t.Error("createHostClient() returned nil")
|
||
return
|
||
}
|
||
|
||
// 检查基本属性
|
||
if client.Addr == "" {
|
||
t.Error("createHostClient() client.Addr is empty")
|
||
}
|
||
|
||
if tt.targetURL == "https://localhost:8443" && !client.IsTLS {
|
||
t.Error("createHostClient() IsTLS should be true for HTTPS")
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// TestHandleWebSocket 测试WebSocket处理
|
||
func TestHandleWebSocket(t *testing.T) {
|
||
cfg := &config.ProxyConfig{
|
||
Path: "/api",
|
||
LoadBalance: "round_robin",
|
||
Timeout: config.ProxyTimeout{Connect: 5 * time.Second},
|
||
}
|
||
|
||
targets := []*loadbalance.Target{
|
||
{URL: "http://localhost:8080", Healthy: true},
|
||
}
|
||
|
||
p, err := NewProxy(cfg, targets)
|
||
if err != nil {
|
||
t.Fatalf("NewProxy() error: %v", err)
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.Header.Set("Upgrade", "websocket")
|
||
ctx.Request.Header.Set("Connection", "upgrade")
|
||
|
||
target := &loadbalance.Target{URL: "http://localhost:8080"}
|
||
client := p.getClient(target.URL)
|
||
|
||
p.handleWebSocket(ctx, target, client)
|
||
|
||
// WebSocket应该返回501 Not Implemented
|
||
if ctx.Response.StatusCode() != fasthttp.StatusNotImplemented {
|
||
t.Errorf("handleWebSocket() status code = %d, want %d", ctx.Response.StatusCode(), fasthttp.StatusNotImplemented)
|
||
}
|
||
}
|
||
|
||
// 辅助函数
|
||
func contains(s, substr string) bool {
|
||
return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsAt(s, substr, 0))
|
||
}
|
||
|
||
func containsAt(s, substr string, start int) bool {
|
||
for i := start; i <= len(s)-len(substr); i++ {
|
||
if s[i:i+len(substr)] == substr {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|