Merge pull request #2641 from tamird/fix-todos

rpc: clean up after tests

rpc/client_test.go

@@ -25,21 +25,9 @@ import (
 	"github.com/cockroachdb/cockroach/util"
 	"github.com/cockroachdb/cockroach/util/hlc"
 	"github.com/cockroachdb/cockroach/util/leaktest"
-	"github.com/cockroachdb/cockroach/util/retry"
 	"github.com/cockroachdb/cockroach/util/stop"
 )
 
-func init() {
-	// Setting these in individual tests triggers the race detector, so
-	// this is better for now.
-	// TODO(tamird): stop doing this.
-	heartbeatInterval = 10 * time.Millisecond
-	clientRetryOptions = retry.Options{
-		InitialBackoff: 1 * time.Millisecond,
-		MaxBackoff:     1 * time.Millisecond,
-	}
-}
-
 func TestClientHeartbeat(t *testing.T) {
 	defer leaktest.AfterTest(t)
 

rpc/main_test.go

@@ -19,12 +19,14 @@ package rpc
 
 import (
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/security"
 	"github.com/cockroachdb/cockroach/security/securitytest"
 	"github.com/cockroachdb/cockroach/testutils"
 	"github.com/cockroachdb/cockroach/util/hlc"
 	"github.com/cockroachdb/cockroach/util/leaktest"
+	"github.com/cockroachdb/cockroach/util/retry"
 	"github.com/cockroachdb/cockroach/util/stop"
 )
 
@@ -44,5 +46,16 @@ func init() {
 //go:generate ../util/leaktest/add-leaktest.sh *_test.go
 
 func TestMain(m *testing.M) {
+	defer func(hbInterval time.Duration, retryOpts retry.Options) {
+		heartbeatInterval = hbInterval
+		clientRetryOptions = retryOpts
+	}(heartbeatInterval, clientRetryOptions)
+
+	heartbeatInterval = 10 * time.Millisecond
+	clientRetryOptions = retry.Options{
+		InitialBackoff: 1 * time.Millisecond,
+		MaxBackoff:     1 * time.Millisecond,
+	}
+
 	leaktest.TestMainWithLeakCheck(m)
 }