From d3905fc2e99fcf97810d7c7301324a4550756aa8 Mon Sep 17 00:00:00 2001
From: yedf2 <120050102@qq.com>
Date: Sat, 18 Dec 2021 08:36:30 +0800
Subject: [PATCH] update tests

---
 app/main.go              |  3 ++-
 bench/main.go            |  3 ++-
 common/config.go         | 12 +++++-------
 common/db.go             | 23 -----------------------
 common/types_test.go     |  4 ----
 common/utils_test.go     |  9 +++++++++
 dtmsvr/storage/boltdb.go |  4 ++++
 dtmsvr/storage/redis.go  |  5 +++++
 dtmsvr/storage/sql.go    | 11 +++++------
 dtmsvr/storage/store.go  |  8 ++++++++
 test/store_test.go       | 11 +++++++++++
 11 files changed, 51 insertions(+), 42 deletions(-)

diff --git a/app/main.go b/app/main.go
index 1cec630..fb97b4b 100644
--- a/app/main.go
+++ b/app/main.go
@@ -15,6 +15,7 @@ import (
 	"github.com/yedf/dtm/dtmcli"
 	"github.com/yedf/dtm/dtmcli/dtmimp"
 	"github.com/yedf/dtm/dtmsvr"
+	"github.com/yedf/dtm/dtmsvr/storage"
 	"github.com/yedf/dtm/examples"
 
 	_ "go.uber.org/automaxprocs"
@@ -53,7 +54,7 @@ func main() {
 	common.MustLoadConfig()
 	dtmcli.SetCurrentDBType(common.Config.ExamplesDB.Driver)
 	if os.Args[1] != "dtmsvr" { // 实际线上运行，只启动dtmsvr，不准备table相关的数据
-		common.WaitDBUp()
+		storage.WaitStoreUp()
 		dtmsvr.PopulateDB(true)
 		examples.PopulateDB(true)
 	}
diff --git a/bench/main.go b/bench/main.go
index 3daf00a..02b2931 100644
--- a/bench/main.go
+++ b/bench/main.go
@@ -8,6 +8,7 @@ import (
 	"github.com/yedf/dtm/dtmcli"
 	"github.com/yedf/dtm/dtmcli/dtmimp"
 	"github.com/yedf/dtm/dtmsvr"
+	"github.com/yedf/dtm/dtmsvr/storage"
 	"github.com/yedf/dtm/examples"
 )
 
@@ -27,7 +28,7 @@ func main() {
 		fmt.Println("start bench server")
 		common.MustLoadConfig()
 		dtmcli.SetCurrentDBType(common.Config.ExamplesDB.Driver)
-		common.WaitDBUp()
+		storage.WaitStoreUp()
 		dtmsvr.PopulateDB(true)
 		examples.PopulateDB(true)
 		dtmsvr.StartSvr()              // 启动dtmsvr的api服务
diff --git a/common/config.go b/common/config.go
index 5f11e6c..7744323 100644
--- a/common/config.go
+++ b/common/config.go
@@ -90,16 +90,14 @@ func MustLoadConfig() {
 }
 
 func checkConfig() error {
-	if Config.Store.Driver == "boltdb" {
-		return nil
-	} else if Config.Store.Driver == "redis" && (Config.Store.Host == "" || Config.Store.Port == 0) {
-		return errors.New("db redis config not valid")
-	} else if Config.Store.Driver != "redis" && (Config.Store.User == "" || Config.Store.Host == "" || Config.Store.Port == 0) {
-		return errors.New("db config not valid")
-	} else if Config.RetryInterval < 10 {
+	if Config.RetryInterval < 10 {
 		return errors.New("RetryInterval should not be less than 10")
 	} else if Config.TimeoutToFail < Config.RetryInterval {
 		return errors.New("TimeoutToFail should not be less than RetryInterval")
+	} else if Config.Store.Driver == "boltdb" {
+		return nil
+	} else if Config.Store.Driver != "redis" && (Config.Store.User == "" || Config.Store.Host == "" || Config.Store.Port == 0) {
+		return errors.New("db config not valid")
 	}
 	return nil
 }
diff --git a/common/db.go b/common/db.go
index 4baf68c..27a07dc 100644
--- a/common/db.go
+++ b/common/db.go
@@ -1,7 +1,6 @@
 package common
 
 import (
-	"context"
 	"database/sql"
 	"fmt"
 	"strings"
@@ -127,25 +126,3 @@ func DbGet(conf dtmcli.DBConf) *DB {
 	}
 	return db.(*DB)
 }
-
-// WaitDBUp wait for db to go up
-func WaitDBUp() {
-	if GetDriver() == "redis" {
-		rdb := RedisGet()
-		for _, err := rdb.Ping(context.Background()).Result(); err != nil; { // wait for mysql to start
-			time.Sleep(3 * time.Second)
-			_, err = rdb.Ping(context.Background()).Result()
-		}
-		return
-	} else if Config.Store.IsDB() {
-		sdb, err := dtmimp.StandaloneDB(Config.Store.GetDBConf())
-		dtmimp.FatalIfError(err)
-		defer func() {
-			sdb.Close()
-		}()
-		for _, err = dtmimp.DBExec(sdb, "select 1"); err != nil; { // wait for mysql to start
-			time.Sleep(3 * time.Second)
-			_, err = dtmimp.DBExec(sdb, "select 1")
-		}
-	}
-}
diff --git a/common/types_test.go b/common/types_test.go
index 4607647..7890c6c 100644
--- a/common/types_test.go
+++ b/common/types_test.go
@@ -32,10 +32,6 @@ func testSql(t *testing.T) {
 	assert.NotEqual(t, nil, err)
 }
 
-func TestWaitDBUp(t *testing.T) {
-	WaitDBUp()
-}
-
 func testDbAlone(t *testing.T) {
 	db, err := dtmimp.StandaloneDB(Config.Store.GetDBConf())
 	assert.Nil(t, err)
diff --git a/common/utils_test.go b/common/utils_test.go
index 851c575..1a9f722 100644
--- a/common/utils_test.go
+++ b/common/utils_test.go
@@ -8,6 +8,7 @@ package common
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -45,3 +46,11 @@ func TestFuncs(t *testing.T) {
 	assert.Equal(t, true, strings.HasSuffix(dir1, "common"))
 
 }
+
+func TestRecoverPanic(t *testing.T) {
+	err := func() (rerr error) {
+		RecoverPanic(&rerr)
+		panic(fmt.Errorf("an error"))
+	}()
+	assert.Equal(t, "an error", err.Error())
+}
diff --git a/dtmsvr/storage/boltdb.go b/dtmsvr/storage/boltdb.go
index 4b3470c..e1e6cc0 100644
--- a/dtmsvr/storage/boltdb.go
+++ b/dtmsvr/storage/boltdb.go
@@ -84,6 +84,10 @@ func tPutIndex(t *bolt.Tx, unix int64, gid string) {
 	dtmimp.E2P(err)
 }
 
+func (s *BoltdbStore) Ping() error {
+	return nil
+}
+
 func (s *BoltdbStore) PopulateData(skipDrop bool) {
 	if !skipDrop {
 		err := boltGet().Update(func(t *bolt.Tx) error {
diff --git a/dtmsvr/storage/redis.go b/dtmsvr/storage/redis.go
index f7931ec..1b225d1 100644
--- a/dtmsvr/storage/redis.go
+++ b/dtmsvr/storage/redis.go
@@ -16,6 +16,11 @@ var ctx context.Context = context.Background()
 type RedisStore struct {
 }
 
+func (s *RedisStore) Ping() error {
+	_, err := redisGet().Ping(ctx).Result()
+	return err
+}
+
 func (s *RedisStore) PopulateData(skipDrop bool) {
 	_, err := redisGet().FlushAll(ctx).Result()
 	dtmimp.PanicIf(err != nil, err)
diff --git a/dtmsvr/storage/sql.go b/dtmsvr/storage/sql.go
index 9858c2d..324eddd 100644
--- a/dtmsvr/storage/sql.go
+++ b/dtmsvr/storage/sql.go
@@ -15,6 +15,11 @@ import (
 type SqlStore struct {
 }
 
+func (s *SqlStore) Ping() error {
+	dbr := dbGet().Exec("select 1")
+	return dbr.Error
+}
+
 func (s *SqlStore) PopulateData(skipDrop bool) {
 	file := fmt.Sprintf("%s/storage.%s.sql", common.GetCallerCodeDir(), config.Store.Driver)
 	common.RunSQLScript(config.Store.GetDBConf(), file, skipDrop)
@@ -126,9 +131,3 @@ func (s *SqlStore) LockOneGlobalTrans(expireIn time.Duration) *TransGlobalStore
 	dbr = db.Must().Where("owner=?", owner).First(global)
 	return global
 }
-
-func lockTransGlobal(db *gorm.DB, gid string, status string) error {
-	g := &TransGlobalStore{}
-	dbr := db.Clauses(clause.Locking{Strength: "UPDATE"}).Model(g).Where("gid=? and status=?", gid, status).First(g)
-	return wrapError(dbr.Error)
-}
diff --git a/dtmsvr/storage/store.go b/dtmsvr/storage/store.go
index 5dfce77..223818b 100644
--- a/dtmsvr/storage/store.go
+++ b/dtmsvr/storage/store.go
@@ -14,6 +14,7 @@ var ErrShouldRetry = errors.New("storage: ShoudRetry")
 var ErrUniqueConflict = errors.New("storage: UniqueKeyConflict")
 
 type Store interface {
+	Ping() error
 	PopulateData(skipDrop bool)
 	FindTransGlobalStore(gid string) *TransGlobalStore
 	ScanTransGlobalStores(position *string, limit int64) []TransGlobalStore
@@ -37,6 +38,13 @@ func GetStore() Store {
 	return stores[config.Store.Driver]
 }
 
+// WaitStoreUp wait for db to go up
+func WaitStoreUp() {
+	for err := GetStore().Ping(); err != nil; err = GetStore().Ping() {
+		time.Sleep(3 * time.Second)
+	}
+}
+
 func wrapError(err error) error {
 	if err == gorm.ErrRecordNotFound || err == redis.Nil {
 		return ErrNotFound
diff --git a/test/store_test.go b/test/store_test.go
index 52665b5..24a2f22 100644
--- a/test/store_test.go
+++ b/test/store_test.go
@@ -85,3 +85,14 @@ func TestStoreLockTrans(t *testing.T) {
 	g2 = s.LockOneGlobalTrans(2 * time.Duration(config.RetryInterval) * time.Second)
 	assert.Nil(t, g2)
 }
+
+func TestStoreWait(t *testing.T) {
+	storage.WaitStoreUp()
+}
+
+func TestUpdateBranchSql(t *testing.T) {
+	if !config.Store.IsDB() {
+		r := storage.GetStore().UpdateBranchesSql(nil, nil)
+		assert.Nil(t, r)
+	}
+}