Refactor: Extract function

suzaku · Jun 17, 2019 · b2339c2 · b2339c2
1 parent c683fb3
commit b2339c2
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 4 deletions.
diff --git a/drainer/config.go b/drainer/config.go
@@ -227,7 +227,7 @@ func (cfg *Config) configFromFile(path string) error {
 // validate checks whether the configuration is valid
 func (cfg *Config) validate() error {
 	// check ListenAddr
-	strictCheck := os.Getenv("BINLOG_TEST") != "1"
+	strictCheck := !util.IsInTestMode()
 	if err := validateAddr(cfg.ListenAddr, strictCheck); err != nil {
 		return errors.Annotate(err, "invalid addr")
 	}

diff --git a/drainer/config_test.go b/drainer/config_test.go
@@ -113,4 +113,4 @@ func (s *validateAddrSuite) TestStrictOrNot(c *C) {
 	c.Assert(err, NotNil)
 	err = validateAddr("http://127.0.0.1:9090", false)
 	c.Assert(err, IsNil)
-}
+}
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -17,6 +17,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"os"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -237,3 +238,9 @@ func AdjustDuration(v *time.Duration, defValue time.Duration) {
 		*v = defValue
 	}
 }
+
+// IsInTestMode returns true if the BINLOG_TEST environment
+// variable is set.
+func IsInTestMode() bool {
+	return os.Getenv("BINLOG_TEST") == "1"
+}
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
@@ -17,6 +17,7 @@ import (
 	"context"
 	"errors"
 	"net"
+	"os"
 	"testing"
 	"time"
 
@@ -48,6 +49,13 @@ func (s dummyStore) CurrentVersion() (kv.Version, error) {
 	return s.ver, s.err
 }
 
+func (s *utilSuite) TestIsInTestMode(c *C) {
+	c.Assert(IsInTestMode(), IsFalse)
+	os.Setenv("BINLOG_TEST", "1")
+	defer os.Unsetenv("BINLOG_TEST")
+	c.Assert(IsInTestMode(), IsTrue)
+}
+
 func (s *utilSuite) TestQueryLatestTsFromPD(c *C) {
 	ds := dummyStore{err: errors.New("test")}
 	ver, err := QueryLatestTsFromPD(ds)

diff --git a/pump/config.go b/pump/config.go
@@ -206,8 +206,7 @@ func (cfg *Config) validate() error {
 		return errors.Errorf("bad AdvertiseAddr host format: %s, %v", urladv.Host, err)
 	}
 	if !util.IsValidateListenHost(host) || host == "0.0.0.0" {
-		isTestEnv := os.Getenv("BINLOG_TEST") == "1"
-		if !(isTestEnv && host == "127.0.0.1") {
+		if !(util.IsInTestMode() && host == "127.0.0.1") {
 			return errors.Errorf("invalid advertiseAddr host: %v", host)
 		}
 	}