mirror of
https://github.com/telanflow/mps.git
synced 2025-09-26 20:41:25 +08:00
103 lines
2.3 KiB
Go
103 lines
2.3 KiB
Go
package mps
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func newTestServer() *httptest.Server {
|
|
return httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
|
|
query := req.URL.Query()
|
|
text := []byte("hello world")
|
|
if query.Get("text") != "" {
|
|
text = []byte(query.Get("text"))
|
|
}
|
|
|
|
rw.Header().Set("Server", "MPS proxy server")
|
|
_, _ = rw.Write(text)
|
|
}))
|
|
}
|
|
|
|
func HttpGet(rawurl string, proxy func(r *http.Request) (*url.URL, error)) (*http.Response, error) {
|
|
req, _ := http.NewRequest(http.MethodGet, rawurl, nil)
|
|
http.DefaultClient.Transport = &http.Transport{
|
|
Proxy: proxy,
|
|
}
|
|
return http.DefaultClient.Do(req)
|
|
}
|
|
|
|
func TestNewHttpProxy(t *testing.T) {
|
|
srv := newTestServer()
|
|
defer srv.Close()
|
|
|
|
proxy := NewHttpProxy()
|
|
proxySrv := httptest.NewServer(proxy)
|
|
defer proxySrv.Close()
|
|
|
|
resp, err := HttpGet(srv.URL, func(r *http.Request) (*url.URL, error) {
|
|
return url.Parse(proxySrv.URL)
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
body, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
|
|
asserts := assert.New(t)
|
|
asserts.Equal(resp.StatusCode, 200, "statusCode should be equal 200")
|
|
asserts.Equal(int64(len(body)), resp.ContentLength)
|
|
}
|
|
|
|
func TestMiddlewareFunc(t *testing.T) {
|
|
// target server
|
|
srv := newTestServer()
|
|
defer srv.Close()
|
|
|
|
// proxy server
|
|
proxy := NewHttpProxy()
|
|
|
|
// use Middleware
|
|
proxy.UseFunc(func(req *http.Request, ctx *Context) (*http.Response, error) {
|
|
resp, err := ctx.Next(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
buf.WriteString("middleware")
|
|
resp.Body = io.NopCloser(&buf)
|
|
|
|
//
|
|
// You have to reset Content-Length, if you change the Body.
|
|
//resp.ContentLength = int64(buf.Len())
|
|
//resp.Header.Set("Content-Length", strconv.Itoa(buf.Len()))
|
|
|
|
return resp, nil
|
|
})
|
|
proxySrv := httptest.NewServer(proxy)
|
|
defer proxySrv.Close()
|
|
|
|
// send request
|
|
resp, err := HttpGet(srv.URL, func(r *http.Request) (*url.URL, error) {
|
|
return url.Parse(proxySrv.URL)
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
body, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
|
|
asserts := assert.New(t)
|
|
asserts.Equal(resp.StatusCode, 200)
|
|
asserts.Equal(int64(len(body)), resp.ContentLength)
|
|
asserts.Equal(string(body), "middleware")
|
|
}
|