diff --git a/Makefile b/Makefile
index 809e4135..be34c1a4 100644
--- a/Makefile
+++ b/Makefile
@@ -113,6 +113,7 @@ generate:
 	@mockgen -source vendor/github.com/deislabs/oras/pkg/auth/client.go -destination pkg/oras/mock/mock_auth.go -package mock
 	@mockgen -source vendor/github.com/containerd/containerd/remotes/resolver.go -destination pkg/oras/mock/mock_resolver.go -package mock
 	@mockgen -source pkg/model/format.go -destination pkg/model/mock/mock_format.go -package mock
+	@mockgen -source pkg/ormb/ormb.go -destination pkg/ormb/mock/mock_ormb.go -package mock
 
 test: generate
 	@go test -race -coverprofile=coverage.out ./...
diff --git a/pkg/ormb/mock/mock.go b/pkg/ormb/mock/mock.go
deleted file mode 100644
index 9ed51f80..00000000
--- a/pkg/ormb/mock/mock.go
+++ /dev/null
@@ -1,148 +0,0 @@
-// Code generated by MockGen. DO NOT EDIT.
-// Source: pkg/oras/interface.go
-
-// Package mock is a generated GoMock package.
-package mock
-
-import (
-	gomock "github.com/golang/mock/gomock"
-	model "github.com/kleveross/ormb/pkg/model"
-	oci "github.com/kleveross/ormb/pkg/oci"
-	reflect "reflect"
-)
-
-// MockInterface is a mock of Interface interface
-type MockInterface struct {
-	ctrl     *gomock.Controller
-	recorder *MockInterfaceMockRecorder
-}
-
-// MockInterfaceMockRecorder is the mock recorder for MockInterface
-type MockInterfaceMockRecorder struct {
-	mock *MockInterface
-}
-
-// NewMockInterface creates a new mock instance
-func NewMockInterface(ctrl *gomock.Controller) *MockInterface {
-	mock := &MockInterface{ctrl: ctrl}
-	mock.recorder = &MockInterfaceMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use
-func (m *MockInterface) EXPECT() *MockInterfaceMockRecorder {
-	return m.recorder
-}
-
-// Login mocks base method
-func (m *MockInterface) Login(hostname, username, password string, insecure bool) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Login", hostname, username, password, insecure)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Login indicates an expected call of Login
-func (mr *MockInterfaceMockRecorder) Login(hostname, username, password, insecure interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Login", reflect.TypeOf((*MockInterface)(nil).Login), hostname, username, password, insecure)
-}
-
-// Logout mocks base method
-func (m *MockInterface) Logout(hostname string) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "Logout", hostname)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Logout indicates an expected call of Logout
-func (mr *MockInterfaceMockRecorder) Logout(hostname interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Logout", reflect.TypeOf((*MockInterface)(nil).Logout), hostname)
-}
-
-// SaveModel mocks base method
-func (m *MockInterface) SaveModel(ch *model.Model, ref *oci.Reference) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "SaveModel", ch, ref)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// SaveModel indicates an expected call of SaveModel
-func (mr *MockInterfaceMockRecorder) SaveModel(ch, ref interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SaveModel", reflect.TypeOf((*MockInterface)(nil).SaveModel), ch, ref)
-}
-
-// PushModel mocks base method
-func (m *MockInterface) PushModel(ref *oci.Reference) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "PushModel", ref)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// PushModel indicates an expected call of PushModel
-func (mr *MockInterfaceMockRecorder) PushModel(ref interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PushModel", reflect.TypeOf((*MockInterface)(nil).PushModel), ref)
-}
-
-// RemoveModel mocks base method
-func (m *MockInterface) RemoveModel(ref *oci.Reference) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "RemoveModel", ref)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// RemoveModel indicates an expected call of RemoveModel
-func (mr *MockInterfaceMockRecorder) RemoveModel(ref interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveModel", reflect.TypeOf((*MockInterface)(nil).RemoveModel), ref)
-}
-
-// PullModel mocks base method
-func (m *MockInterface) PullModel(ref *oci.Reference) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "PullModel", ref)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// PullModel indicates an expected call of PullModel
-func (mr *MockInterfaceMockRecorder) PullModel(ref interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PullModel", reflect.TypeOf((*MockInterface)(nil).PullModel), ref)
-}
-
-// LoadModel mocks base method
-func (m *MockInterface) LoadModel(ref *oci.Reference) (*model.Model, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "LoadModel", ref)
-	ret0, _ := ret[0].(*model.Model)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// LoadModel indicates an expected call of LoadModel
-func (mr *MockInterfaceMockRecorder) LoadModel(ref interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LoadModel", reflect.TypeOf((*MockInterface)(nil).LoadModel), ref)
-}
-
-// TagModel mocks base method
-func (m *MockInterface) TagModel(ref, target *oci.Reference) error {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "TagModel", ref, target)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// TagModel indicates an expected call of TagModel
-func (mr *MockInterfaceMockRecorder) TagModel(ref, target interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TagModel", reflect.TypeOf((*MockInterface)(nil).TagModel), ref, target)
-}
diff --git a/pkg/ormb/mock/mock_ormb.go b/pkg/ormb/mock/mock_ormb.go
new file mode 100644
index 00000000..cb5ef95a
--- /dev/null
+++ b/pkg/ormb/mock/mock_ormb.go
@@ -0,0 +1,131 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: pkg/ormb/ormb.go
+
+// Package mock is a generated GoMock package.
+package mock
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	reflect "reflect"
+)
+
+// MockInterface is a mock of Interface interface
+type MockInterface struct {
+	ctrl     *gomock.Controller
+	recorder *MockInterfaceMockRecorder
+}
+
+// MockInterfaceMockRecorder is the mock recorder for MockInterface
+type MockInterfaceMockRecorder struct {
+	mock *MockInterface
+}
+
+// NewMockInterface creates a new mock instance
+func NewMockInterface(ctrl *gomock.Controller) *MockInterface {
+	mock := &MockInterface{ctrl: ctrl}
+	mock.recorder = &MockInterfaceMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockInterface) EXPECT() *MockInterfaceMockRecorder {
+	return m.recorder
+}
+
+// Login mocks base method
+func (m *MockInterface) Login(hostname, username, password string, insecureOpt bool) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Login", hostname, username, password, insecureOpt)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Login indicates an expected call of Login
+func (mr *MockInterfaceMockRecorder) Login(hostname, username, password, insecureOpt interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Login", reflect.TypeOf((*MockInterface)(nil).Login), hostname, username, password, insecureOpt)
+}
+
+// Push mocks base method
+func (m *MockInterface) Push(refStr string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Push", refStr)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Push indicates an expected call of Push
+func (mr *MockInterfaceMockRecorder) Push(refStr interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Push", reflect.TypeOf((*MockInterface)(nil).Push), refStr)
+}
+
+// Pull mocks base method
+func (m *MockInterface) Pull(refStr string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Pull", refStr)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Pull indicates an expected call of Pull
+func (mr *MockInterfaceMockRecorder) Pull(refStr interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Pull", reflect.TypeOf((*MockInterface)(nil).Pull), refStr)
+}
+
+// Export mocks base method
+func (m *MockInterface) Export(refStr, dst string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Export", refStr, dst)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Export indicates an expected call of Export
+func (mr *MockInterfaceMockRecorder) Export(refStr, dst interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Export", reflect.TypeOf((*MockInterface)(nil).Export), refStr, dst)
+}
+
+// Save mocks base method
+func (m *MockInterface) Save(src, refStr string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Save", src, refStr)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Save indicates an expected call of Save
+func (mr *MockInterfaceMockRecorder) Save(src, refStr interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Save", reflect.TypeOf((*MockInterface)(nil).Save), src, refStr)
+}
+
+// Tag mocks base method
+func (m *MockInterface) Tag(refStr, targetStr string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Tag", refStr, targetStr)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Tag indicates an expected call of Tag
+func (mr *MockInterfaceMockRecorder) Tag(refStr, targetStr interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Tag", reflect.TypeOf((*MockInterface)(nil).Tag), refStr, targetStr)
+}
+
+// Remove mocks base method
+func (m *MockInterface) Remove(refStr string) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Remove", refStr)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Remove indicates an expected call of Remove
+func (mr *MockInterfaceMockRecorder) Remove(refStr interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Remove", reflect.TypeOf((*MockInterface)(nil).Remove), refStr)
+}