2019-04-22 02:59:20 +00:00

171 lines
4.3 KiB
Go

package proxy
import (
"bytes"
"context"
"net/http"
"net/url"
"sync"
"testing"
"time"
"go-common/library/log"
bm "go-common/library/net/http/blademaster"
"github.com/stretchr/testify/assert"
)
func init() {
log.Init(nil)
}
func TestProxy(t *testing.T) {
engine := bm.Default()
engine.GET("/icon", NewAlways("http://api.bilibili.com/x/web-interface/index/icon"))
engine.POST("/x/web-interface/archive/like", NewAlways("http://api.bilibili.com"))
go engine.Run(":18080")
defer func() {
engine.Server().Shutdown(context.TODO())
}()
time.Sleep(time.Second)
req, err := http.NewRequest("GET", "http://127.0.0.1:18080/icon", nil)
assert.NoError(t, err)
req.Host = "api.bilibili.com"
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
// proxy form request
form := url.Values{}
form.Set("arg1", "1")
form.Set("arg2", "2")
req, err = http.NewRequest("POST", "http://127.0.0.1:18080/x/web-interface/archive/like?param=test", bytes.NewReader([]byte(form.Encode())))
assert.NoError(t, err)
req.Host = "api.bilibili.com"
resp, err = http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
// proxy json request
bs := []byte(`{"arg1": 1, "arg2": 2}`)
req, err = http.NewRequest("POST", "http://127.0.0.1:18080/x/web-interface/archive/like?param=test", bytes.NewReader(bs))
assert.NoError(t, err)
req.Host = "api.bilibili.com"
req.Header.Set("Content-Type", "application/json; charset=utf-8")
resp, err = http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
}
func TestProxyRace(t *testing.T) {
engine := bm.Default()
engine.GET("/icon", NewAlways("http://api.bilibili.com/x/web-interface/index/icon"))
go engine.Run(":18080")
defer func() {
engine.Server().Shutdown(context.TODO())
}()
time.Sleep(time.Second)
wg := sync.WaitGroup{}
for i := 0; i < 20; i++ {
wg.Add(1)
go func() {
defer wg.Done()
req, err := http.NewRequest("GET", "http://127.0.0.1:18080/icon", nil)
assert.NoError(t, err)
req.Host = "api.bilibili.com"
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
}()
}
wg.Wait()
}
func TestZoneProxy(t *testing.T) {
engine := bm.Default()
engine.GET("/icon", NewZoneProxy("sh004", "http://api.bilibili.com/x/web-interface/index/icon"), func(ctx *bm.Context) {
ctx.AbortWithStatus(500)
})
engine.GET("/icon2", NewZoneProxy("none", "http://api.bilibili.com/x/web-interface/index/icon2"), func(ctx *bm.Context) {
ctx.AbortWithStatus(200)
})
ug := engine.Group("/update", NewZoneProxy("sh004", "http://api.bilibili.com"))
ug.POST("/name", func(ctx *bm.Context) {
ctx.AbortWithStatus(500)
})
ug.POST("/sign", func(ctx *bm.Context) {
ctx.AbortWithStatus(500)
})
go engine.Run(":18080")
defer func() {
engine.Server().Shutdown(context.TODO())
}()
time.Sleep(time.Second)
req, err := http.NewRequest("GET", "http://127.0.0.1:18080/icon", nil)
assert.NoError(t, err)
req.Host = "api.bilibili.com"
req.Header.Set("X-BILI-SLB", "shjd-out-slb")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
req.URL.Path = "/icon2"
resp, err = http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
req.URL.Path = "/update/name"
resp, err = http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
req.URL.Path = "/update/sign"
resp, err = http.DefaultClient.Do(req)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, 200, resp.StatusCode)
}
func BenchmarkProxy(b *testing.B) {
engine := bm.Default()
engine.GET("/icon", NewAlways("http://api.bilibili.com/x/web-interface/index/icon"))
go engine.Run(":18080")
defer func() {
engine.Server().Shutdown(context.TODO())
}()
time.Sleep(time.Second)
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
req, err := http.NewRequest("GET", "http://127.0.0.1:18080/icon", nil)
assert.NoError(b, err)
req.Host = "api.bilibili.com"
resp, err := http.DefaultClient.Do(req)
assert.NoError(b, err)
defer resp.Body.Close()
assert.Equal(b, 200, resp.StatusCode)
}
})
}