You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
137 lines
3.0 KiB
137 lines
3.0 KiB
package blademaster
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
criticalityPkg "github.com/bilibili/kratos/pkg/net/criticality"
|
|
"github.com/bilibili/kratos/pkg/net/metadata"
|
|
xtime "github.com/bilibili/kratos/pkg/time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
var (
|
|
sonce sync.Once
|
|
|
|
SockAddr = "localhost:18080"
|
|
curEngine atomic.Value
|
|
)
|
|
|
|
func uri(base, path string) string {
|
|
return fmt.Sprintf("%s://%s%s", "http", base, path)
|
|
}
|
|
|
|
func shutdown() {
|
|
if err := curEngine.Load().(*Engine).Shutdown(context.TODO()); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func setupHandler(engine *Engine) {
|
|
// set the global timeout is 2 second
|
|
engine.conf.Timeout = xtime.Duration(time.Second * 2)
|
|
|
|
engine.Ping(func(ctx *Context) {
|
|
ctx.AbortWithStatus(200)
|
|
})
|
|
|
|
engine.GET("/criticality/api", Criticality(criticalityPkg.Critical), func(ctx *Context) {
|
|
ctx.String(200, "%s", metadata.String(ctx, metadata.Criticality))
|
|
})
|
|
engine.GET("/criticality/none/api", func(ctx *Context) {
|
|
ctx.String(200, "%s", metadata.String(ctx, metadata.Criticality))
|
|
})
|
|
}
|
|
|
|
func startServer() {
|
|
e := DefaultServer(nil)
|
|
setupHandler(e)
|
|
go e.Run(SockAddr)
|
|
curEngine.Store(e)
|
|
time.Sleep(time.Second)
|
|
}
|
|
|
|
func TestCriticality(t *testing.T) {
|
|
startServer()
|
|
defer shutdown()
|
|
|
|
tests := []*struct {
|
|
path string
|
|
crtl criticalityPkg.Criticality
|
|
expected criticalityPkg.Criticality
|
|
}{
|
|
{
|
|
"/criticality/api",
|
|
criticalityPkg.EmptyCriticality,
|
|
criticalityPkg.Critical,
|
|
},
|
|
{
|
|
"/criticality/api",
|
|
criticalityPkg.CriticalPlus,
|
|
criticalityPkg.Critical,
|
|
},
|
|
{
|
|
"/criticality/api",
|
|
criticalityPkg.SheddablePlus,
|
|
criticalityPkg.Critical,
|
|
},
|
|
}
|
|
client := &http.Client{}
|
|
for _, testCase := range tests {
|
|
req, err := http.NewRequest("GET", uri(SockAddr, testCase.path), nil)
|
|
assert.NoError(t, err)
|
|
req.Header.Set(_httpHeaderCriticality, string(testCase.crtl))
|
|
resp, err := client.Do(req)
|
|
assert.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
body, err := ioutil.ReadAll(resp.Body)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, testCase.expected, criticalityPkg.Criticality(body))
|
|
}
|
|
}
|
|
|
|
func TestNoneCriticality(t *testing.T) {
|
|
startServer()
|
|
defer shutdown()
|
|
|
|
tests := []*struct {
|
|
path string
|
|
crtl criticalityPkg.Criticality
|
|
expected criticalityPkg.Criticality
|
|
}{
|
|
{
|
|
"/criticality/none/api",
|
|
criticalityPkg.EmptyCriticality,
|
|
criticalityPkg.Critical,
|
|
},
|
|
{
|
|
"/criticality/none/api",
|
|
criticalityPkg.CriticalPlus,
|
|
criticalityPkg.CriticalPlus,
|
|
},
|
|
{
|
|
"/criticality/none/api",
|
|
criticalityPkg.SheddablePlus,
|
|
criticalityPkg.SheddablePlus,
|
|
},
|
|
}
|
|
client := &http.Client{}
|
|
for _, testCase := range tests {
|
|
req, err := http.NewRequest("GET", uri(SockAddr, testCase.path), nil)
|
|
assert.NoError(t, err)
|
|
req.Header.Set(_httpHeaderCriticality, string(testCase.crtl))
|
|
resp, err := client.Do(req)
|
|
assert.NoError(t, err)
|
|
defer resp.Body.Close()
|
|
body, err := ioutil.ReadAll(resp.Body)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, testCase.expected, criticalityPkg.Criticality(body))
|
|
}
|
|
}
|
|
|