120 lines
3.6 KiB
Go
120 lines
3.6 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
)
|
|
|
|
type ProxyProbeServiceSuite struct {
|
|
suite.Suite
|
|
ctx context.Context
|
|
proxySrv *httptest.Server
|
|
prober *proxyProbeService
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) SetupTest() {
|
|
s.ctx = context.Background()
|
|
s.prober = &proxyProbeService{
|
|
ipInfoURL: "http://ip-api.test/json/?lang=zh-CN",
|
|
allowPrivateHosts: true,
|
|
}
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TearDownTest() {
|
|
if s.proxySrv != nil {
|
|
s.proxySrv.Close()
|
|
s.proxySrv = nil
|
|
}
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) setupProxyServer(handler http.HandlerFunc) {
|
|
s.proxySrv = newLocalTestServer(s.T(), handler)
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_InvalidProxyURL() {
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, "://bad")
|
|
require.Error(s.T(), err)
|
|
require.ErrorContains(s.T(), err, "failed to create proxy client")
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_UnsupportedProxyScheme() {
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, "ftp://127.0.0.1:1")
|
|
require.Error(s.T(), err)
|
|
require.ErrorContains(s.T(), err, "failed to create proxy client")
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_Success() {
|
|
seen := make(chan string, 1)
|
|
s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
seen <- r.RequestURI
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{"status":"success","query":"1.2.3.4","city":"c","regionName":"r","country":"cc","countryCode":"CC"}`)
|
|
}))
|
|
|
|
info, latencyMs, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
|
|
require.NoError(s.T(), err, "ProbeProxy")
|
|
require.GreaterOrEqual(s.T(), latencyMs, int64(0), "unexpected latency")
|
|
require.Equal(s.T(), "1.2.3.4", info.IP)
|
|
require.Equal(s.T(), "c", info.City)
|
|
require.Equal(s.T(), "r", info.Region)
|
|
require.Equal(s.T(), "cc", info.Country)
|
|
require.Equal(s.T(), "CC", info.CountryCode)
|
|
|
|
// Verify proxy received the request
|
|
select {
|
|
case uri := <-seen:
|
|
require.Contains(s.T(), uri, "ip-api.test", "expected request to go through proxy")
|
|
default:
|
|
require.Fail(s.T(), "expected proxy to receive request")
|
|
}
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_NonOKStatus() {
|
|
s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
}))
|
|
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
|
|
require.Error(s.T(), err)
|
|
require.ErrorContains(s.T(), err, "status: 503")
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_InvalidJSON() {
|
|
s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, "not-json")
|
|
}))
|
|
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
|
|
require.Error(s.T(), err)
|
|
require.ErrorContains(s.T(), err, "failed to parse response")
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_InvalidIPInfoURL() {
|
|
s.prober.ipInfoURL = "://invalid-url"
|
|
s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
|
|
require.Error(s.T(), err, "expected error for invalid ipInfoURL")
|
|
}
|
|
|
|
func (s *ProxyProbeServiceSuite) TestProbeProxy_ProxyServerClosed() {
|
|
s.setupProxyServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
|
s.proxySrv.Close()
|
|
|
|
_, _, err := s.prober.ProbeProxy(s.ctx, s.proxySrv.URL)
|
|
require.Error(s.T(), err, "expected error when proxy server is closed")
|
|
}
|
|
|
|
func TestProxyProbeServiceSuite(t *testing.T) {
|
|
suite.Run(t, new(ProxyProbeServiceSuite))
|
|
}
|