Add SelectOne() to VizierDBInterface

Signed-off-by: Koichiro Den <den@valinux.co.jp>
kubeflow · Dec 4, 2018 · 61ac560 · 61ac560
1 parent 65c3406
commit 61ac560
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/pkg/db/db_init.go b/pkg/db/db_init.go
@@ -1,6 +1,7 @@
 package db
 
 import (
+	"fmt"
 	"log"
 )
 
@@ -94,3 +95,12 @@ func (d *dbConn) DBInit() {
 		log.Fatalf("Error creating earlystop_param table: %v", err)
 	}
 }
+
+func (d *dbConn) SelectOne() error {
+	db := d.db
+	_, err := db.Exec(`SELECT 1`)
+	if err != nil {
+		return fmt.Errorf("Error `SELECT 1` probing: %v", err)
+	}
+	return nil
+}
diff --git a/pkg/db/interface.go b/pkg/db/interface.go
@@ -44,6 +44,8 @@ type WorkerLog struct {
 
 type VizierDBInterface interface {
 	DBInit()
+	SelectOne() error
+
 	GetStudyConfig(string) (*api.StudyConfig, error)
 	GetStudyList() ([]string, error)
 	CreateStudy(*api.StudyConfig) (string, error)

diff --git a/pkg/db/interface_test.go b/pkg/db/interface_test.go
@@ -12,7 +12,6 @@ import (
 
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/golang/protobuf/jsonpb"
-	"gopkg.in/DATA-DOG/go-sqlmock.v1"
 
 	api "github.com/kubeflow/katib/pkg/api"
 )
@@ -48,6 +47,10 @@ func TestMain(m *testing.M) {
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS suggestion_param").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS earlystop_param").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	dbInterface.DBInit()
+	err = dbInterface.SelectOne()
+	if err != nil {
+		fmt.Printf("error `SELECT 1` probing: %v\n", err)
+	}
 
 	mysqlAddr := os.Getenv("TEST_MYSQL")
 	if mysqlAddr != "" {