fix stats test race (#1560)

* fix stats test race
* and try fix TestMapAddressEnv race

proxy_test.go

@@ -46,29 +46,6 @@ func overwrite(hpfe func(req *http.Request) (*url.URL, error)) func() {
 	}
 }
 
-func TestMapAddressEnv(t *testing.T) {
-	// Overwrite the function in the test and restore them in defer.
-	hpfe := func(req *http.Request) (*url.URL, error) {
-		if req.URL.Host == envTestAddr {
-			return &url.URL{
-				Scheme: "https",
-				Host:   envProxyAddr,
-			}, nil
-		}
-		return nil, nil
-	}
-	defer overwrite(hpfe)()
-
-	// envTestAddr should be handled by ProxyFromEnvironment.
-	got, err := mapAddress(context.Background(), envTestAddr)
-	if err != nil {
-		t.Error(err)
-	}
-	if got != envProxyAddr {
-		t.Errorf("want %v, got %v", envProxyAddr, got)
-	}
-}
-
 type proxyServer struct {
 	t   *testing.T
 	lis net.Listener
@@ -177,3 +154,27 @@ func TestHTTPConnect(t *testing.T) {
 		t.Fatalf("received msg: %v, want %v", recvBuf, msg)
 	}
 }
+
+func TestMapAddressEnv(t *testing.T) {
+	defer leakcheck.Check(t)
+	// Overwrite the function in the test and restore them in defer.
+	hpfe := func(req *http.Request) (*url.URL, error) {
+		if req.URL.Host == envTestAddr {
+			return &url.URL{
+				Scheme: "https",
+				Host:   envProxyAddr,
+			}, nil
+		}
+		return nil, nil
+	}
+	defer overwrite(hpfe)()
+
+	// envTestAddr should be handled by ProxyFromEnvironment.
+	got, err := mapAddress(context.Background(), envTestAddr)
+	if err != nil {
+		t.Error(err)
+	}
+	if got != envProxyAddr {
+		t.Errorf("want %v, got %v", envProxyAddr, got)
+	}
+}

stats/stats_test.go

@@ -1108,7 +1108,9 @@ func testClientStats(t *testing.T, tc *testConfig, cc *rpcConfig, checkFuncs map
 		err:         err,
 	}
 
+	h.mu.Lock()
 	checkConnStats(t, h.gotConn)
+	h.mu.Unlock()
 	checkClientStats(t, h.gotRPC, expect, checkFuncs)
 }