diff --git a/db/interface_test.go b/db/interface_test.go
index ddd71a2c80d..91de1eaeec8 100644
--- a/db/interface_test.go
+++ b/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()
 
diff --git a/manager/main_test.go b/manager/main_test.go
index 01385cdaf07..71bc785f4cb 100644
--- a/manager/main_test.go
+++ b/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{
diff --git a/mock/mock_api/manager_mock.go b/mock/api/manager.go
similarity index 59%
rename from mock/mock_api/manager_mock.go
rename to mock/api/manager.go
index 9fee8953e02..9c90e0e0efd 100644
--- a/mock/mock_api/manager_mock.go
+++ b/mock/api/manager.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/kubeflow/hp-tuning/api (interfaces: ManagerClient)
 
-// Package mock_api is a generated GoMock package.
-package mock_api
+// Package mock is a generated GoMock package.
+package mock
 
 import (
 	context "context"
@@ -89,6 +89,24 @@ func (mr *MockManagerClientMockRecorder) CreateStudy(arg0, arg1 interface{}, arg
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateStudy", reflect.TypeOf((*MockManagerClient)(nil).CreateStudy), varargs...)
 }
 
+// EarlyStopping mocks base method
+func (m *MockManagerClient) EarlyStopping(arg0 context.Context, arg1 *api.EarlyStoppingRequest, arg2 ...grpc.CallOption) (*api.EarlyStoppingReply, error) {
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "EarlyStopping", varargs...)
+	ret0, _ := ret[0].(*api.EarlyStoppingReply)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// EarlyStopping indicates an expected call of EarlyStopping
+func (mr *MockManagerClientMockRecorder) EarlyStopping(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EarlyStopping", reflect.TypeOf((*MockManagerClient)(nil).EarlyStopping), varargs...)
+}
+
 // GetObjectValue mocks base method
 func (m *MockManagerClient) GetObjectValue(arg0 context.Context, arg1 *api.GetObjectValueRequest, arg2 ...grpc.CallOption) (*api.GetObjectValueReply, error) {
 	varargs := []interface{}{arg0, arg1}
@@ -107,22 +125,76 @@ func (mr *MockManagerClientMockRecorder) GetObjectValue(arg0, arg1 interface{},
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetObjectValue", reflect.TypeOf((*MockManagerClient)(nil).GetObjectValue), varargs...)
 }
 
-// GetStudys mocks base method
-func (m *MockManagerClient) GetStudys(arg0 context.Context, arg1 *api.GetStudysRequest, arg2 ...grpc.CallOption) (*api.GetStudysReply, error) {
+// GetSavedModel mocks base method
+func (m *MockManagerClient) GetSavedModel(arg0 context.Context, arg1 *api.GetSavedModelRequest, arg2 ...grpc.CallOption) (*api.GetSavedModelReply, error) {
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetSavedModel", varargs...)
+	ret0, _ := ret[0].(*api.GetSavedModelReply)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedModel indicates an expected call of GetSavedModel
+func (mr *MockManagerClientMockRecorder) GetSavedModel(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedModel", reflect.TypeOf((*MockManagerClient)(nil).GetSavedModel), varargs...)
+}
+
+// GetSavedModels mocks base method
+func (m *MockManagerClient) GetSavedModels(arg0 context.Context, arg1 *api.GetSavedModelsRequest, arg2 ...grpc.CallOption) (*api.GetSavedModelsReply, error) {
 	varargs := []interface{}{arg0, arg1}
 	for _, a := range arg2 {
 		varargs = append(varargs, a)
 	}
-	ret := m.ctrl.Call(m, "GetStudys", varargs...)
-	ret0, _ := ret[0].(*api.GetStudysReply)
+	ret := m.ctrl.Call(m, "GetSavedModels", varargs...)
+	ret0, _ := ret[0].(*api.GetSavedModelsReply)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetStudys indicates an expected call of GetStudys
-func (mr *MockManagerClientMockRecorder) GetStudys(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+// GetSavedModels indicates an expected call of GetSavedModels
+func (mr *MockManagerClientMockRecorder) GetSavedModels(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
 	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetStudys", reflect.TypeOf((*MockManagerClient)(nil).GetStudys), varargs...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedModels", reflect.TypeOf((*MockManagerClient)(nil).GetSavedModels), varargs...)
+}
+
+// GetSavedStudies mocks base method
+func (m *MockManagerClient) GetSavedStudies(arg0 context.Context, arg1 *api.GetSavedStudiesRequest, arg2 ...grpc.CallOption) (*api.GetSavedStudiesReply, error) {
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetSavedStudies", varargs...)
+	ret0, _ := ret[0].(*api.GetSavedStudiesReply)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedStudies indicates an expected call of GetSavedStudies
+func (mr *MockManagerClientMockRecorder) GetSavedStudies(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedStudies", reflect.TypeOf((*MockManagerClient)(nil).GetSavedStudies), varargs...)
+}
+
+// GetStudies mocks base method
+func (m *MockManagerClient) GetStudies(arg0 context.Context, arg1 *api.GetStudiesRequest, arg2 ...grpc.CallOption) (*api.GetStudiesReply, error) {
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "GetStudies", varargs...)
+	ret0, _ := ret[0].(*api.GetStudiesReply)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetStudies indicates an expected call of GetStudies
+func (mr *MockManagerClientMockRecorder) GetStudies(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetStudies", reflect.TypeOf((*MockManagerClient)(nil).GetStudies), varargs...)
 }
 
 // InitializeSuggestService mocks base method
@@ -143,22 +215,40 @@ func (mr *MockManagerClientMockRecorder) InitializeSuggestService(arg0, arg1 int
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InitializeSuggestService", reflect.TypeOf((*MockManagerClient)(nil).InitializeSuggestService), varargs...)
 }
 
-// ShouldTrialStop mocks base method
-func (m *MockManagerClient) ShouldTrialStop(arg0 context.Context, arg1 *api.ShouldTrialStopRequest, arg2 ...grpc.CallOption) (*api.ShouldTrialStopReply, error) {
+// SaveModel mocks base method
+func (m *MockManagerClient) SaveModel(arg0 context.Context, arg1 *api.SaveModelRequest, arg2 ...grpc.CallOption) (*api.SaveModelReply, error) {
+	varargs := []interface{}{arg0, arg1}
+	for _, a := range arg2 {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "SaveModel", varargs...)
+	ret0, _ := ret[0].(*api.SaveModelReply)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// SaveModel indicates an expected call of SaveModel
+func (mr *MockManagerClientMockRecorder) SaveModel(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{arg0, arg1}, arg2...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveModel", reflect.TypeOf((*MockManagerClient)(nil).SaveModel), varargs...)
+}
+
+// SaveStudy mocks base method
+func (m *MockManagerClient) SaveStudy(arg0 context.Context, arg1 *api.SaveStudyRequest, arg2 ...grpc.CallOption) (*api.SaveStudyReply, error) {
 	varargs := []interface{}{arg0, arg1}
 	for _, a := range arg2 {
 		varargs = append(varargs, a)
 	}
-	ret := m.ctrl.Call(m, "ShouldTrialStop", varargs...)
-	ret0, _ := ret[0].(*api.ShouldTrialStopReply)
+	ret := m.ctrl.Call(m, "SaveStudy", varargs...)
+	ret0, _ := ret[0].(*api.SaveStudyReply)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// ShouldTrialStop indicates an expected call of ShouldTrialStop
-func (mr *MockManagerClientMockRecorder) ShouldTrialStop(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
+// SaveStudy indicates an expected call of SaveStudy
+func (mr *MockManagerClientMockRecorder) SaveStudy(arg0, arg1 interface{}, arg2 ...interface{}) *gomock.Call {
 	varargs := append([]interface{}{arg0, arg1}, arg2...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ShouldTrialStop", reflect.TypeOf((*MockManagerClient)(nil).ShouldTrialStop), varargs...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveStudy", reflect.TypeOf((*MockManagerClient)(nil).SaveStudy), varargs...)
 }
 
 // StopStudy mocks base method
diff --git a/mock/mock_api/suggestion_mock.go b/mock/api/suggestion.go
similarity index 98%
rename from mock/mock_api/suggestion_mock.go
rename to mock/api/suggestion.go
index e82dfff405a..9a108e0e91c 100644
--- a/mock/mock_api/suggestion_mock.go
+++ b/mock/api/suggestion.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/kubeflow/hp-tuning/api (interfaces: SuggestionClient)
 
-// Package mock_api is a generated GoMock package.
-package mock_api
+// Package mock is a generated GoMock package.
+package mock
 
 import (
 	context "context"
diff --git a/mock/mock_db/db_mock.go b/mock/db/db.go
similarity index 99%
rename from mock/mock_db/db_mock.go
rename to mock/db/db.go
index 7cf6857389f..09b1efb7c0a 100644
--- a/mock/mock_db/db_mock.go
+++ b/mock/db/db.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/kubeflow/hp-tuning/db (interfaces: VizierDBInterface)
 
-// Package mock_db is a generated GoMock package.
-package mock_db
+// Package mock is a generated GoMock package.
+package mock
 
 import (
 	gomock "github.com/golang/mock/gomock"
diff --git a/mock/modelstore/modelstore.go b/mock/modelstore/modelstore.go
new file mode 100644
index 00000000000..1fe6bfd375b
--- /dev/null
+++ b/mock/modelstore/modelstore.go
@@ -0,0 +1,97 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/kubeflow/hp-tuning/manager/modelstore (interfaces: ModelStore)
+
+// Package mock is a generated GoMock package.
+package mock
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	api "github.com/kubeflow/hp-tuning/api"
+	reflect "reflect"
+)
+
+// MockModelStore is a mock of ModelStore interface
+type MockModelStore struct {
+	ctrl     *gomock.Controller
+	recorder *MockModelStoreMockRecorder
+}
+
+// MockModelStoreMockRecorder is the mock recorder for MockModelStore
+type MockModelStoreMockRecorder struct {
+	mock *MockModelStore
+}
+
+// NewMockModelStore creates a new mock instance
+func NewMockModelStore(ctrl *gomock.Controller) *MockModelStore {
+	mock := &MockModelStore{ctrl: ctrl}
+	mock.recorder = &MockModelStoreMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockModelStore) EXPECT() *MockModelStoreMockRecorder {
+	return m.recorder
+}
+
+// GetSavedModel mocks base method
+func (m *MockModelStore) GetSavedModel(arg0 *api.GetSavedModelRequest) (*api.ModelInfo, error) {
+	ret := m.ctrl.Call(m, "GetSavedModel", arg0)
+	ret0, _ := ret[0].(*api.ModelInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedModel indicates an expected call of GetSavedModel
+func (mr *MockModelStoreMockRecorder) GetSavedModel(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedModel", reflect.TypeOf((*MockModelStore)(nil).GetSavedModel), arg0)
+}
+
+// GetSavedModels mocks base method
+func (m *MockModelStore) GetSavedModels(arg0 *api.GetSavedModelsRequest) ([]*api.ModelInfo, error) {
+	ret := m.ctrl.Call(m, "GetSavedModels", arg0)
+	ret0, _ := ret[0].([]*api.ModelInfo)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedModels indicates an expected call of GetSavedModels
+func (mr *MockModelStoreMockRecorder) GetSavedModels(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedModels", reflect.TypeOf((*MockModelStore)(nil).GetSavedModels), arg0)
+}
+
+// GetSavedStudies mocks base method
+func (m *MockModelStore) GetSavedStudies() ([]*api.StudyOverview, error) {
+	ret := m.ctrl.Call(m, "GetSavedStudies")
+	ret0, _ := ret[0].([]*api.StudyOverview)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetSavedStudies indicates an expected call of GetSavedStudies
+func (mr *MockModelStoreMockRecorder) GetSavedStudies() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSavedStudies", reflect.TypeOf((*MockModelStore)(nil).GetSavedStudies))
+}
+
+// SaveModel mocks base method
+func (m *MockModelStore) SaveModel(arg0 *api.SaveModelRequest) error {
+	ret := m.ctrl.Call(m, "SaveModel", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SaveModel indicates an expected call of SaveModel
+func (mr *MockModelStoreMockRecorder) SaveModel(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveModel", reflect.TypeOf((*MockModelStore)(nil).SaveModel), arg0)
+}
+
+// SaveStudy mocks base method
+func (m *MockModelStore) SaveStudy(arg0 *api.SaveStudyRequest) error {
+	ret := m.ctrl.Call(m, "SaveStudy", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// SaveStudy indicates an expected call of SaveStudy
+func (mr *MockModelStoreMockRecorder) SaveStudy(arg0 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveStudy", reflect.TypeOf((*MockModelStore)(nil).SaveStudy), arg0)
+}
diff --git a/mock/mock_worker_interface/worker_interface_mock.go b/mock/worker/worker.go
similarity index 87%
rename from mock/mock_worker_interface/worker_interface_mock.go
rename to mock/worker/worker.go
index 262d5d5ab07..a8eb2624d33 100644
--- a/mock/mock_worker_interface/worker_interface_mock.go
+++ b/mock/worker/worker.go
@@ -1,8 +1,8 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: github.com/kubeflow/hp-tuning/manager/worker_interface (interfaces: WorkerInterface)
 
-// Package mock_worker_interface is a generated GoMock package.
-package mock_worker_interface
+// Package mock is a generated GoMock package.
+package mock
 
 import (
 	gomock "github.com/golang/mock/gomock"
@@ -34,15 +34,15 @@ func (m *MockWorkerInterface) EXPECT() *MockWorkerInterfaceMockRecorder {
 }
 
 // CheckRunningTrials mocks base method
-func (m *MockWorkerInterface) CheckRunningTrials(arg0, arg1 string, arg2 []string) error {
-	ret := m.ctrl.Call(m, "CheckRunningTrials", arg0, arg1, arg2)
+func (m *MockWorkerInterface) CheckRunningTrials(arg0, arg1 string) error {
+	ret := m.ctrl.Call(m, "CheckRunningTrials", arg0, arg1)
 	ret0, _ := ret[0].(error)
 	return ret0
 }
 
 // CheckRunningTrials indicates an expected call of CheckRunningTrials
-func (mr *MockWorkerInterfaceMockRecorder) CheckRunningTrials(arg0, arg1, arg2 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRunningTrials", reflect.TypeOf((*MockWorkerInterface)(nil).CheckRunningTrials), arg0, arg1, arg2)
+func (mr *MockWorkerInterfaceMockRecorder) CheckRunningTrials(arg0, arg1 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CheckRunningTrials", reflect.TypeOf((*MockWorkerInterface)(nil).CheckRunningTrials), arg0, arg1)
 }
 
 // CleanWorkers mocks base method
@@ -57,6 +57,18 @@ func (mr *MockWorkerInterfaceMockRecorder) CleanWorkers(arg0 interface{}) *gomoc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CleanWorkers", reflect.TypeOf((*MockWorkerInterface)(nil).CleanWorkers), arg0)
 }
 
+// CompleteTrial mocks base method
+func (m *MockWorkerInterface) CompleteTrial(arg0, arg1 string, arg2 bool) error {
+	ret := m.ctrl.Call(m, "CompleteTrial", arg0, arg1, arg2)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// CompleteTrial indicates an expected call of CompleteTrial
+func (mr *MockWorkerInterfaceMockRecorder) CompleteTrial(arg0, arg1, arg2 interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CompleteTrial", reflect.TypeOf((*MockWorkerInterface)(nil).CompleteTrial), arg0, arg1, arg2)
+}
+
 // GetCompletedTrials mocks base method
 func (m *MockWorkerInterface) GetCompletedTrials(arg0 string) []*api.Trial {
 	ret := m.ctrl.Call(m, "GetCompletedTrials", arg0)
diff --git a/scripts/mockgen.sh b/scripts/mockgen.sh
new file mode 100755
index 00000000000..fffaa444c02
--- /dev/null
+++ b/scripts/mockgen.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# Copyright 2018 The Kubeflow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+SCRIPT_ROOT=$(dirname ${BASH_SOURCE})/..
+
+cd ${SCRIPT_ROOT}
+echo "Generating ManagerClient..."
+mockgen -package mock -destination mock/api/manager.go github.com/kubeflow/hp-tuning/api ManagerClient
+echo "Generating SuggestionClient..."
+mockgen -package mock -destination mock/api/suggestion.go github.com/kubeflow/hp-tuning/api SuggestionClient
+echo "Generating VizierDBInterface..."
+mockgen -package mock -destination mock/db/db.go github.com/kubeflow/hp-tuning/db VizierDBInterface
+echo "Generating WorkerInterface..."
+mockgen -package mock -destination mock/worker/worker.go  github.com/kubeflow/hp-tuning/manager/worker_interface WorkerInterface
+echo "Generating ModelStore..."
+mockgen -package mock -destination mock/modelstore/modelstore.go  github.com/kubeflow/hp-tuning/manager/modelstore ModelStore
+cd - > /dev/null
diff --git a/test/scripts/build.sh b/test/scripts/build.sh
index 0befb5891e5..e3ef76606e0 100755
--- a/test/scripts/build.sh
+++ b/test/scripts/build.sh
@@ -28,9 +28,11 @@ VERSION=$(git describe --tags --always --dirty)
 
 echo "Activating service-account"
 gcloud auth activate-service-account --key-file=${GOOGLE_APPLICATION_CREDENTIALS}
+
 echo "Create symlink to GOPATH"
 mkdir -p ${GOPATH}/src/github.com/${REPO_OWNER}
 ln -s ${PWD} ${GO_DIR}
+
 cd ${GO_DIR}
 echo "Build operator binary"
 mkdir bin
diff --git a/test/scripts/unit-test.sh b/test/scripts/unit-test.sh
new file mode 100755
index 00000000000..74da54ca1f0
--- /dev/null
+++ b/test/scripts/unit-test.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+
+# Copyright 2018 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This shell script is used to build a cluster and create a namespace from our
+# argo workflow
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+export PATH=${GOPATH}/bin:/usr/local/go/bin:${PATH}
+REGISTRY="${GCP_REGISTRY}"
+PROJECT="${GCP_PROJECT}"
+GO_DIR=${GOPATH}/src/github.com/${REPO_OWNER}/${REPO_NAME}
+VERSION=$(git describe --tags --always --dirty)
+
+echo "Activating service-account"
+gcloud auth activate-service-account --key-file=${GOOGLE_APPLICATION_CREDENTIALS}
+
+echo "Run unit test cases"
+cd ${GO_DIR}
+go test ./...
+cd - > /dev/null
diff --git a/test/workflows/components/workflows.libsonnet b/test/workflows/components/workflows.libsonnet
index c5eed48b79f..3b87b638256 100644
--- a/test/workflows/components/workflows.libsonnet
+++ b/test/workflows/components/workflows.libsonnet
@@ -191,7 +191,7 @@
               },
             },
           ],  // volumes
-          // onExit specifies the template that should always run when the workflow completes.          
+          // onExit specifies the template that should always run when the workflow completes.
           onExit: "exit-handler",
           templates: [
             {
@@ -211,6 +211,12 @@
                     template: "create-pr-symlink",
                   },
                 ],
+                [
+                  {
+                    name: "unit-test",
+                    template: "unit-test",
+                  },
+                ],
                 [  // Setup cluster needs to run after build because we depend on the chart
                   // created by the build statement.
                   {
@@ -288,6 +294,9 @@
             $.parts(namespace, name).e2e(prow_env, bucket).buildTemplate("build", testWorkerImage, [
               "test/scripts/build.sh",
             ]),  // build
+            $.parts(namespace, name).e2e(prow_env, bucket).buildTemplate("unit-test", testWorkerImage, [
+              "test/scripts/unit-test.sh",
+            ]),  // unit test
           ],  // templates
         },
       },  // e2e