From b2339c2531b31ebc946b7deb6a2955ee68924e8c Mon Sep 17 00:00:00 2001
From: satoru <satorulogic@gmail.com>
Date: Mon, 17 Jun 2019 11:59:38 +0800
Subject: [PATCH] Refactor: Extract function

---
 drainer/config.go      | 2 +-
 drainer/config_test.go | 2 +-
 pkg/util/util.go       | 7 +++++++
 pkg/util/util_test.go  | 8 ++++++++
 pump/config.go         | 3 +--
 5 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/drainer/config.go b/drainer/config.go
index 1378685d5..9bc817994 100644
--- 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
index 0530f7c06..6dbe28463 100644
--- 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)
-}
\ No newline at end of file
+}
diff --git a/pkg/util/util.go b/pkg/util/util.go
index cd886e8b9..c6c65dc5a 100644
--- 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
index f21ce5dd2..476ba7acf 100644
--- 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
index b2b28fc2b..c029e3df0 100644
--- 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)
 		}
 	}