Merge pull request #62 from proxymo-network/custom-client-support

Add speedtest struct which contains the global variables used
showwin · Feb 20, 2022 · 4e10a07 · 4e10a07
2 parents a83af84 + cf8e53c
commit 4e10a07
Show file tree

Hide file tree

Showing 8 changed files with 121 additions and 30 deletions.
diff --git a/speedtest/request.go b/speedtest/request.go
@@ -13,14 +13,13 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-type downloadWarmUpFunc func(context.Context, string) error
-type downloadFunc func(context.Context, string, int) error
-type uploadWarmUpFunc func(context.Context, string) error
-type uploadFunc func(context.Context, string, int) error
+type downloadWarmUpFunc func(context.Context, *http.Client, string) error
+type downloadFunc func(context.Context, *http.Client, string, int) error
+type uploadWarmUpFunc func(context.Context, *http.Client, string) error
+type uploadFunc func(context.Context, *http.Client, string, int) error
 
 var dlSizes = [...]int{350, 500, 750, 1000, 1500, 2000, 2500, 3000, 3500, 4000}
 var ulSizes = [...]int{100, 300, 500, 800, 1000, 1500, 2500, 3000, 3500, 4000} //kB
-var client = http.Client{}
 
 // DownloadTest executes the test to measure download speed
 func (s *Server) DownloadTest(savingMode bool) error {
@@ -45,7 +44,7 @@ func (s *Server) downloadTestContext(
 	sTime := time.Now()
 	for i := 0; i < 2; i++ {
 		eg.Go(func() error {
-			return dlWarmUp(ctx, dlURL)
+			return dlWarmUp(ctx, s.doer, dlURL)
 		})
 	}
 	if err := eg.Wait(); err != nil {
@@ -84,7 +83,7 @@ func (s *Server) downloadTestContext(
 		sTime = time.Now()
 		for i := 0; i < workload; i++ {
 			eg.Go(func() error {
-				return downloadRequest(ctx, dlURL, weight)
+				return downloadRequest(ctx, s.doer, dlURL, weight)
 			})
 		}
 		if err := eg.Wait(); err != nil {
@@ -120,7 +119,7 @@ func (s *Server) uploadTestContext(
 	eg := errgroup.Group{}
 	for i := 0; i < 2; i++ {
 		eg.Go(func() error {
-			return ulWarmUp(ctx, s.URL)
+			return ulWarmUp(ctx, s.doer, s.URL)
 		})
 	}
 	if err := eg.Wait(); err != nil {
@@ -159,7 +158,7 @@ func (s *Server) uploadTestContext(
 		sTime = time.Now()
 		for i := 0; i < workload; i++ {
 			eg.Go(func() error {
-				return uploadRequest(ctx, s.URL, weight)
+				return uploadRequest(ctx, s.doer, s.URL, weight)
 			})
 		}
 		if err := eg.Wait(); err != nil {
@@ -176,7 +175,7 @@ func (s *Server) uploadTestContext(
 	return nil
 }
 
-func dlWarmUp(ctx context.Context, dlURL string) error {
+func dlWarmUp(ctx context.Context, doer *http.Client, dlURL string) error {
 	size := dlSizes[2]
 	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"
 
@@ -185,7 +184,7 @@ func dlWarmUp(ctx context.Context, dlURL string) error {
 		return err
 	}
 
-	resp, err := client.Do(req)
+	resp, err := doer.Do(req)
 	if err != nil {
 		return err
 	}
@@ -194,7 +193,7 @@ func dlWarmUp(ctx context.Context, dlURL string) error {
 	return err
 }
 
-func ulWarmUp(ctx context.Context, ulURL string) error {
+func ulWarmUp(ctx context.Context, doer *http.Client, ulURL string) error {
 	size := ulSizes[4]
 	v := url.Values{}
 	v.Add("content", strings.Repeat("0123456789", size*100-51))
@@ -205,7 +204,7 @@ func ulWarmUp(ctx context.Context, ulURL string) error {
 	}
 
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-	resp, err := client.Do(req)
+	resp, err := doer.Do(req)
 	if err != nil {
 		return err
 	}
@@ -214,7 +213,7 @@ func ulWarmUp(ctx context.Context, ulURL string) error {
 	return err
 }
 
-func downloadRequest(ctx context.Context, dlURL string, w int) error {
+func downloadRequest(ctx context.Context, doer *http.Client, dlURL string, w int) error {
 	size := dlSizes[w]
 	xdlURL := dlURL + "/random" + strconv.Itoa(size) + "x" + strconv.Itoa(size) + ".jpg"
 
@@ -223,7 +222,7 @@ func downloadRequest(ctx context.Context, dlURL string, w int) error {
 		return err
 	}
 
-	resp, err := client.Do(req)
+	resp, err := doer.Do(req)
 	if err != nil {
 		return err
 	}
@@ -232,7 +231,7 @@ func downloadRequest(ctx context.Context, dlURL string, w int) error {
 	return err
 }
 
-func uploadRequest(ctx context.Context, ulURL string, w int) error {
+func uploadRequest(ctx context.Context, doer *http.Client, ulURL string, w int) error {
 	size := ulSizes[w]
 	v := url.Values{}
 	v.Add("content", strings.Repeat("0123456789", size*100-51))
@@ -243,7 +242,7 @@ func uploadRequest(ctx context.Context, ulURL string, w int) error {
 	}
 
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-	resp, err := client.Do(req)
+	resp, err := doer.Do(req)
 	if err != nil {
 		return err
 	}
@@ -262,7 +261,7 @@ func (s *Server) PingTest() error {
 func (s *Server) PingTestContext(ctx context.Context) error {
 	pingURL := strings.Split(s.URL, "/upload.php")[0] + "/latency.txt"
 
-	l := time.Duration(100000000000) // 10sec
+	l := time.Second * 10
 	for i := 0; i < 3; i++ {
 		sTime := time.Now()
 
@@ -271,7 +270,7 @@ func (s *Server) PingTestContext(ctx context.Context) error {
 			return err
 		}
 
-		resp, err := http.DefaultClient.Do(req)
+		resp, err := s.doer.Do(req)
 		if err != nil {
 			return err
 		}

diff --git a/speedtest/request_test.go b/speedtest/request_test.go
@@ -3,6 +3,7 @@ package speedtest
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"testing"
 	"time"
 )
@@ -91,12 +92,12 @@ func TestUploadTestContextSavingMode(t *testing.T) {
 	}
 }
 
-func mockWarmUp(ctx context.Context, dlURL string) error {
+func mockWarmUp(ctx context.Context, doer *http.Client, dlURL string) error {
 	time.Sleep(100 * time.Millisecond)
 	return nil
 }
 
-func mockRequest(ctx context.Context, dlURL string, w int) error {
+func mockRequest(ctx context.Context, doer *http.Client, dlURL string, w int) error {
 	fmt.Sprintln(w)
 	time.Sleep(500 * time.Millisecond)
 	return nil

diff --git a/speedtest/server.go b/speedtest/server.go
@@ -30,6 +30,8 @@ type Server struct {
 	Latency  time.Duration `json:"latency"`
 	DLSpeed  float64       `json:"dl_speed"`
 	ULSpeed  float64       `json:"ul_speed"`
+
+	doer *http.Client
 }
 
 // ServerList list of Server
@@ -60,19 +62,24 @@ func (b ByDistance) Less(i, j int) bool {
 	return b.Servers[i].Distance < b.Servers[j].Distance
 }
 
+// FetchServerList retrieves a list of available servers
+func (client *Speedtest) FetchServerList(user *User) (ServerList, error) {
+	return client.FetchServerListContext(context.Background(), user)
+}
+
 // FetchServerList retrieves a list of available servers
 func FetchServerList(user *User) (ServerList, error) {
-	return FetchServerListContext(context.Background(), user)
+	return defaultClient.FetchServerList(user)
 }
 
 // FetchServerListContext retrieves a list of available servers, observing the given context.
-func FetchServerListContext(ctx context.Context, user *User) (ServerList, error) {
+func (client *Speedtest) FetchServerListContext(ctx context.Context, user *User) (ServerList, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, speedTestServersUrl, nil)
 	if err != nil {
 		return ServerList{}, err
 	}
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := client.doer.Do(req)
 	if err != nil {
 		return ServerList{}, err
 	}
@@ -85,7 +92,7 @@ func FetchServerListContext(ctx context.Context, user *User) (ServerList, error)
 			return ServerList{}, err
 		}
 
-		resp, err = http.DefaultClient.Do(req)
+		resp, err = client.doer.Do(req)
 		if err != nil {
 			return ServerList{}, err
 		}
@@ -101,6 +108,11 @@ func FetchServerListContext(ctx context.Context, user *User) (ServerList, error)
 		return list, err
 	}
 
+	// set doer of server
+	for _, s := range list.Servers {
+		s.doer = client.doer
+	}
+
 	// Calculate distance
 	for i := range list.Servers {
 		server := list.Servers[i]
@@ -121,6 +133,11 @@ func FetchServerListContext(ctx context.Context, user *User) (ServerList, error)
 	return list, nil
 }
 
+// FetchServerListContext retrieves a list of available servers, observing the given context.
+func FetchServerListContext(ctx context.Context, user *User) (ServerList, error) {
+	return defaultClient.FetchServerListContext(ctx, user)
+}
+
 func distance(lat1 float64, lon1 float64, lat2 float64, lon2 float64) float64 {
 	radius := 6378.137
 

diff --git a/speedtest/server_test.go b/speedtest/server_test.go
@@ -9,7 +9,10 @@ func TestFetchServerList(t *testing.T) {
 		Lon: "138.44",
 		Isp: "Hello",
 	}
-	serverList, err := FetchServerList(&user)
+
+	client := New()
+
+	serverList, err := client.FetchServerList(&user)
 	if err != nil {
 		t.Errorf(err.Error())
 	}

diff --git a/speedtest/speedtest.go b/speedtest/speedtest.go
@@ -0,0 +1,33 @@
+package speedtest
+
+import "net/http"
+
+// Speedtest is a speedtest client.
+type Speedtest struct {
+	doer *http.Client
+}
+
+// Option is a function that can be passed to New to modify the Client.
+type Option func(*Speedtest)
+
+// WithDoer sets the http.Client used to make requests.
+func WithDoer(doer *http.Client) Option {
+	return func(s *Speedtest) {
+		s.doer = doer
+	}
+}
+
+// New creates a new speedtest client.
+func New(opts ...Option) *Speedtest {
+	s := &Speedtest{
+		doer: http.DefaultClient,
+	}
+
+	for _, opt := range opts {
+		opt(s)
+	}
+
+	return s
+}
+
+var defaultClient = New()
diff --git a/speedtest/speedtest_test.go b/speedtest/speedtest_test.go
@@ -0,0 +1,26 @@
+package speedtest
+
+import (
+	"net/http"
+	"testing"
+)
+
+func TestNew(t *testing.T) {
+	t.Run("DefaultDoer", func(t *testing.T) {
+		c := New()
+
+		if c.doer == nil {
+			t.Error("doer is nil by")
+		}
+	})
+
+	t.Run("CustomDoer", func(t *testing.T) {
+		doer := &http.Client{}
+
+		c := New(WithDoer(doer))
+		if c.doer != doer {
+			t.Error("doer is not the same")
+		}
+	})
+
+}
diff --git a/speedtest/user.go b/speedtest/user.go
@@ -23,19 +23,24 @@ type Users struct {
 	Users []User `xml:"client"`
 }
 
+// FetchUserInfo returns information about caller determined by speedtest.net
+func (client *Speedtest) FetchUserInfo() (*User, error) {
+	return client.FetchUserInfoContext(context.Background())
+}
+
 // FetchUserInfo returns information about caller determined by speedtest.net
 func FetchUserInfo() (*User, error) {
-	return FetchUserInfoContext(context.Background())
+	return defaultClient.FetchUserInfo()
 }
 
 // FetchUserInfoContext returns information about caller determined by speedtest.net, observing the given context.
-func FetchUserInfoContext(ctx context.Context) (*User, error) {
+func (client *Speedtest) FetchUserInfoContext(ctx context.Context) (*User, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, speedTestConfigUrl, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := http.DefaultClient.Do(req)
+	resp, err := client.doer.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -57,6 +62,11 @@ func FetchUserInfoContext(ctx context.Context) (*User, error) {
 	return &users.Users[0], nil
 }
 
+// FetchUserInfoContext returns information about caller determined by speedtest.net, observing the given context.
+func FetchUserInfoContext(ctx context.Context) (*User, error) {
+	return defaultClient.FetchUserInfoContext(ctx)
+}
+
 // String representation of User
 func (u *User) String() string {
 	return fmt.Sprintf("%s, (%s) [%s, %s]", u.IP, u.Isp, u.Lat, u.Lon)

diff --git a/speedtest/user_test.go b/speedtest/user_test.go
@@ -7,7 +7,9 @@ import (
 )
 
 func TestFetchUserInfo(t *testing.T) {
-	user, err := FetchUserInfo()
+	client := New()
+
+	user, err := client.FetchUserInfo()
 	if err != nil {
 		t.Errorf(err.Error())
 	}