argo: Add unit test

db/interface_test.go

@@ -6,15 +6,15 @@ package db
 
 import (
 	"fmt"
-	"github.com/golang/protobuf/jsonpb"
 	"math/rand"
 	"os"
 	"testing"
 
-	api "github.com/kubeflow/hp-tuning/api"
-
 	_ "github.com/go-sql-driver/mysql"
+	"github.com/golang/protobuf/jsonpb"
 	"gopkg.in/DATA-DOG/go-sqlmock.v1"
+
+	api "github.com/kubeflow/hp-tuning/api"
 )
 
 var db_interface VizierDBInterface
@@ -33,7 +33,8 @@ func TestMain(m *testing.M) {
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS studies").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS study_permissions").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS trials").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
-	mock.ExpectExec("CREATE TABLE IF NOT EXISTS trial_logs").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectExec("CREATE TABLE IF NOT EXISTS trial_metrics").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectExec("CREATE TABLE IF NOT EXISTS trial_lastlogs").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectExec("CREATE TABLE IF NOT EXISTS workers").WithArgs().WillReturnResult(sqlmock.NewResult(1, 1))
 	db_interface.DB_Init()
 

manager/main_test.go

@@ -2,19 +2,22 @@ package main
 
 import (
 	"context"
+	"testing"
+
 	"github.com/golang/mock/gomock"
+
 	api "github.com/kubeflow/hp-tuning/api"
-	//"github.com/kubeflow/hp-tuning/mock/mock_api"
-	"github.com/kubeflow/hp-tuning/mock/mock_db"
-	"github.com/kubeflow/hp-tuning/mock/mock_worker_interface"
-	"testing"
+	mockdb "github.com/kubeflow/hp-tuning/mock/db"
+	mockmodelstore "github.com/kubeflow/hp-tuning/mock/modelstore"
+	mockworker "github.com/kubeflow/hp-tuning/mock/worker"
 )
 
 func TestCreateStudy(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
-	mockDB := mock_db.NewMockVizierDBInterface(ctrl)
-	mockWif := mock_worker_interface.NewMockWorkerInterface(ctrl)
+	mockDB := mockdb.NewMockVizierDBInterface(ctrl)
+	mockWif := mockworker.NewMockWorkerInterface(ctrl)
+	mockModelStore := mockmodelstore.NewMockModelStore(ctrl)
 	sid := "teststudy"
 	sc := &api.StudyConfig{
 		Name:               "test",
@@ -27,7 +30,16 @@ func TestCreateStudy(t *testing.T) {
 	mockDB.EXPECT().CreateStudy(
 		sc,
 	).Return(sid, nil)
-	s := &server{wIF: mockWif, StudyChList: make(map[string]studyCh)}
+	ssr := &api.SaveStudyRequest{
+		StudyName: "test",
+	}
+	mockModelStore.EXPECT().SaveStudy(ssr).Return(nil)
+
+	s := &server{
+		wIF:         mockWif,
+		msIf:        mockModelStore,
+		StudyChList: make(map[string]studyCh),
+	}
 	req := &api.CreateStudyRequest{StudyConfig: sc}
 	ret, err := s.CreateStudy(context.Background(), req)
 	if err != nil {
@@ -48,10 +60,18 @@ func TestCreateStudy(t *testing.T) {
 func TestGetStudies(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
-	mockDB := mock_db.NewMockVizierDBInterface(ctrl)
-	mockWif := mock_worker_interface.NewMockWorkerInterface(ctrl)
+	mockDB := mockdb.NewMockVizierDBInterface(ctrl)
+	mockWif := mockworker.NewMockWorkerInterface(ctrl)
+	mockModelStore := mockmodelstore.NewMockModelStore(ctrl)
 	sid := []string{"teststudy1", "teststudy2"}
-	s := &server{wIF: mockWif, StudyChList: map[string]studyCh{sid[0]: studyCh{}, sid[1]: studyCh{}}}
+	s := &server{
+		wIF:  mockWif,
+		msIf: mockModelStore,
+		StudyChList: map[string]studyCh{
+			sid[0]: studyCh{},
+			sid[1]: studyCh{},
+		},
+	}
 	dbIf = mockDB
 
 	sc := []*api.StudyConfig{