diff --git a/pkg/controller.v1beta1/experiment/manifest/generator_test.go b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
index 0a35c93f052..895510847f4 100644
--- a/pkg/controller.v1beta1/experiment/manifest/generator_test.go
+++ b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
@@ -22,28 +22,31 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
-	"go.uber.org/mock/gomock"
 	batchv1 "k8s.io/api/batch/v1"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	commonapiv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
 	experimentsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/experiments/v1beta1"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/util"
-	katibclientmock "github.com/kubeflow/katib/pkg/mock/v1beta1/util/katibclient"
+	"github.com/kubeflow/katib/pkg/util/v1beta1/katibclient"
 )
 
 func TestGetRunSpecWithHP(t *testing.T) {
 
-	mockCtrl := gomock.NewController(t)
-	defer mockCtrl.Finish()
+	scheme := runtime.NewScheme()
+	v1.AddToScheme(scheme)
+	fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
 
-	c := katibclientmock.NewMockClient(mockCtrl)
+	// Use the same pattern as in the New() function in generator.go
+	katibClient := katibclient.NewWithGivenClient(fakeClient)
 
 	p := &DefaultGenerator{
-		client: c,
+		client: katibClient,
 	}
 
 	expectedJob := batchv1.Job{
@@ -152,15 +155,7 @@ func TestGetRunSpecWithHP(t *testing.T) {
 }
 
 func TestGetRunSpecWithHPConfigMap(t *testing.T) {
-	mockCtrl := gomock.NewController(t)
-	defer mockCtrl.Finish()
-
-	c := katibclientmock.NewMockClient(mockCtrl)
-
-	p := &DefaultGenerator{
-		client: c,
-	}
-
+	// Mocking the ConfigMap
 	templatePath := "trial-template-path"
 
 	trialSpec := `apiVersion: batch/v1
@@ -222,17 +217,23 @@ spec:
 	}
 
 	cases := map[string]struct {
-		mockConfigMapGetter            func() *gomock.Call
+		objects                        []runtime.Object
 		instance                       *experimentsv1beta1.Experiment
 		parameterAssignments           []commonapiv1beta1.ParameterAssignment
 		wantRunSpecWithHyperParameters *unstructured.Unstructured
 		wantError                      error
 	}{
 		"Run with valid parameters": {
-			mockConfigMapGetter: func() *gomock.Call {
-				return c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-					map[string]string{templatePath: trialSpec}, nil,
-				)
+			objects: []runtime.Object{
+				&v1.ConfigMap{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "config-map-name",
+						Namespace: "config-map-namespace",
+					},
+					Data: map[string]string{
+						templatePath: trialSpec,
+					},
+				},
 			},
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -249,10 +250,16 @@ spec:
 			wantRunSpecWithHyperParameters: expectedRunSpec,
 		},
 		"Invalid ConfigMap name": {
-			mockConfigMapGetter: func() *gomock.Call {
-				return c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-					nil, errConfigMapNotFound,
-				)
+			objects: []runtime.Object{
+				&v1.ConfigMap{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "config-map-name",
+						Namespace: "config-map-namespace",
+					},
+					Data: map[string]string{
+						templatePath: trialSpec,
+					},
+				},
 			},
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -267,10 +274,16 @@ spec:
 			wantError:            errConfigMapNotFound,
 		},
 		"Invalid template path in ConfigMap name": {
-			mockConfigMapGetter: func() *gomock.Call {
-				return c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-					map[string]string{templatePath: trialSpec}, nil,
-				)
+			objects: []runtime.Object{
+				&v1.ConfigMap{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "config-map-name",
+						Namespace: "config-map-namespace",
+					},
+					Data: map[string]string{
+						templatePath: trialSpec, //No templatePath
+					},
+				},
 			},
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -289,10 +302,16 @@ spec:
 		// Trial template is a string in ConfigMap
 		// Because of that, user can specify not valid unstructured template
 		"Invalid trial spec in ConfigMap": {
-			mockConfigMapGetter: func() *gomock.Call {
-				return c.EXPECT().GetConfigMap(gomock.Any(), gomock.Any()).Return(
-					map[string]string{templatePath: invalidTrialSpec}, nil,
-				)
+			objects: []runtime.Object{
+				&v1.ConfigMap{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "config-map-name",
+						Namespace: "config-map-namespace",
+					},
+					Data: map[string]string{
+						templatePath: invalidTrialSpec,
+					},
+				},
 			},
 			instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -312,7 +331,16 @@ spec:
 
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {
-			tc.mockConfigMapGetter()
+			scheme := runtime.NewScheme()
+			v1.AddToScheme(scheme)
+			fakeClient := fake.NewClientBuilder().WithScheme(scheme).WithRuntimeObjects(tc.objects...).Build()
+
+			// Use the same pattern as in the New() function in generator.go
+			katibClient := katibclient.NewWithGivenClient(fakeClient)
+
+			p := &DefaultGenerator{
+				client: katibClient,
+			}
 			got, err := p.GetRunSpecWithHyperParameters(tc.instance, "trial-name", "trial-namespace", tc.parameterAssignments)
 			if diff := cmp.Diff(tc.wantError, err, cmpopts.EquateErrors()); len(diff) != 0 {
 				t.Errorf("Unexpected error from GetRunSpecWithHyperParameters (-want,+got):\n%s", diff)