kratos/pkg/net/http/blademaster/server_test.go

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("x-bm-metadata-criticality", 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("x-bm-metadata-criticality", 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))
	}
}