@@ -2,6 +2,7 @@ package gateway
22
33import (
44 "io"
5+ "math/rand"
56 "net"
67 "net/http"
78 "net/http/httptest"
@@ -11,12 +12,17 @@ import (
1112)
1213
1314type mockDNSResolver struct {
14- IPs []net.IP
15- Err error
15+ IPs []net.IP
16+ Err error
17+ Rand * rand.Rand
1618}
1719
1820func (m mockDNSResolver ) LookupIP (hostname string ) ([]net.IP , error ) {
19- return m .IPs , m .Err
21+ shuffledIPs := make ([]net.IP , len (m .IPs ))
22+ copy (shuffledIPs , m .IPs )
23+ rand .Shuffle (len (shuffledIPs ), func (i , j int ) { shuffledIPs [i ], shuffledIPs [j ] = shuffledIPs [j ], shuffledIPs [i ] })
24+
25+ return shuffledIPs , m .Err
2026}
2127
2228type customRoundTripper struct {}
@@ -31,6 +37,7 @@ func (rt customRoundTripper) RoundTrip(req *http.Request) (*http.Response, error
3137}
3238
3339func TestDistribution (t * testing.T ) {
40+ r := rand .New (rand .NewSource (time .Now ().UnixNano ()))
3441 hostname := "example.com"
3542 testCases := []struct {
3643 name string
@@ -109,8 +116,9 @@ func TestDistribution(t *testing.T) {
109116 for _ , tc := range testCases {
110117 t .Run (tc .name , func (t * testing.T ) {
111118 mockResolver := mockDNSResolver {
112- IPs : tc .IPs ,
113- Err : nil ,
119+ IPs : tc .IPs ,
120+ Err : nil ,
121+ Rand : r ,
114122 }
115123
116124 lb , err := newRoundRobinLoadBalancer (hostname , mockResolver .LookupIP )
0 commit comments