Storage: Extract returned object hydration into function (#88012)

* Extract returned object hydration into function

* Finish writing tests for utils func

* Lint

* Update pkg/apiserver/rest/dualwriter_mode2.go

Co-authored-by: Arati R. <33031346+suntala@users.noreply.github.com>

* Better var naming

* Remove duplicated logic

* Lint

* Fix test

* Lint

* Make type private

* Fix one more test

* Fix test

---------

Co-authored-by: Arati R. <33031346+suntala@users.noreply.github.com>

pkg/apiserver/rest/dualwriter_mode1.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"time"
 
-	"k8s.io/apimachinery/pkg/api/meta"
 	metainternalversion "k8s.io/apimachinery/pkg/apis/meta/internalversion"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -16,8 +15,8 @@ import (
 type DualWriterMode1 struct {
 	Legacy  LegacyStorage
 	Storage Storage
-	Log     klog.Logger
 	*dualWriterMetrics
+	Log klog.Logger
 }
 
 const (
@@ -38,40 +37,34 @@ func (d *DualWriterMode1) Mode() DualWriterMode {
 }
 
 // Create overrides the behavior of the generic DualWriter and writes only to LegacyStorage.
-func (d *DualWriterMode1) Create(ctx context.Context, obj runtime.Object, createValidation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
+func (d *DualWriterMode1) Create(ctx context.Context, original runtime.Object, createValidation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
 	log := d.Log.WithValues("kind", options.Kind)
 	ctx = klog.NewContext(ctx, log)
 	var method = "create"
 
 	startLegacy := time.Now()
-	res, err := d.Legacy.Create(ctx, obj, createValidation, options)
+	created, err := d.Legacy.Create(ctx, original, createValidation, options)
 	if err != nil {
 		log.Error(err, "unable to create object in legacy storage")
 		d.recordLegacyDuration(true, mode1Str, options.Kind, method, startLegacy)
-		return res, err
+		return created, err
 	}
 	d.recordLegacyDuration(false, mode1Str, options.Kind, method, startLegacy)
 
 	go func() {
-		accessorCreated, err := meta.Accessor(res)
-		if err != nil {
-			log.Error(err, "unable to get accessor for created object")
-		}
-
-		accessorOld, err := meta.Accessor(obj)
-		if err != nil {
-			log.Error(err, "unable to get accessor for old object")
-		}
-
-		enrichObject(accessorOld, accessorCreated)
-		startStorage := time.Now()
 		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage create timeout"))
+		createdLegacy, err := enrichLegacyObject(original, created, true)
+		if err != nil {
+			cancel()
+		}
+
+		startStorage := time.Now()
 		defer cancel()
-		_, errObjectSt := d.Storage.Create(ctx, obj, createValidation, options)
+		_, errObjectSt := d.Storage.Create(ctx, createdLegacy, createValidation, options)
 		d.recordStorageDuration(errObjectSt != nil, mode1Str, options.Kind, method, startStorage)
 	}()
 
-	return res, nil
+	return created, nil
 }
 
 // Get overrides the behavior of the generic DualWriter and reads only from LegacyStorage.
@@ -188,6 +181,7 @@ func (d *DualWriterMode1) Update(ctx context.Context, name string, objInfo rest.
 	d.recordLegacyDuration(false, mode1Str, options.Kind, method, startLegacy)
 
 	go func() {
+		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage update timeout"))
 		updated, err := objInfo.UpdatedObject(ctx, res)
 		if err != nil {
 			log.WithValues("object", updated).Error(err, "could not update or create object")
@@ -201,27 +195,17 @@ func (d *DualWriterMode1) Update(ctx context.Context, name string, objInfo rest.
 
 		// if the object is found, create a new updateWrapper with the object found
 		if foundObj != nil {
-			accessorOld, err := meta.Accessor(foundObj)
+			res, err := enrichLegacyObject(foundObj, res, false)
 			if err != nil {
-				log.Error(err, "unable to get accessor for original updated object")
+				log.Error(err, "could not enrich object")
+				cancel()
 			}
-
-			accessor, err := meta.Accessor(res)
-			if err != nil {
-				log.Error(err, "unable to get accessor for updated object")
-			}
-
-			accessor.SetResourceVersion(accessorOld.GetResourceVersion())
-			accessor.SetUID(accessorOld.GetUID())
-
-			enrichObject(accessorOld, accessor)
 			objInfo = &updateWrapper{
 				upstream: objInfo,
 				updated:  res,
 			}
 		}
 		startStorage := time.Now()
-		ctx, cancel := context.WithTimeoutCause(ctx, time.Second*10, errors.New("storage update timeout"))
 		defer cancel()
 		_, _, errObjectSt := d.Storage.Update(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
 		d.recordStorageDuration(errObjectSt != nil, mode1Str, options.Kind, method, startStorage)

pkg/apiserver/rest/dualwriter_mode1_test.go

@@ -22,10 +22,10 @@ var anotherList = &example.PodList{Items: []example.Pod{*anotherObj}}
 
 func TestMode1_Create(t *testing.T) {
 	type testCase struct {
-		name           string
 		input          runtime.Object
 		setupLegacyFn  func(m *mock.Mock, input runtime.Object)
 		setupStorageFn func(m *mock.Mock, input runtime.Object)
+		name           string
 		wantErr        bool
 	}
 	tests :=
@@ -51,6 +51,7 @@ func TestMode1_Create(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -71,22 +72,23 @@ func TestMode1_Create(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			us.AssertNotCalled(t, "Create", context.Background(), tt.input, func(context.Context, runtime.Object) error { return nil }, &metav1.CreateOptions{})
 
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode1_Get(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock, name string)
 		setupStorageFn func(m *mock.Mock, name string)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -112,6 +114,7 @@ func TestMode1_Get(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -132,21 +135,22 @@ func TestMode1_Get(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			us.AssertNotCalled(t, "Get", context.Background(), tt.name, &metav1.GetOptions{})
 
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode1_List(t *testing.T) {
 	type testCase struct {
-		name           string
 		setupLegacyFn  func(m *mock.Mock)
 		setupStorageFn func(m *mock.Mock)
+		name           string
 		wantErr        bool
 	}
 	tests :=
@@ -161,6 +165,7 @@ func TestMode1_List(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -181,17 +186,18 @@ func TestMode1_List(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
+		})
 	}
 }
 
 func TestMode1_Delete(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock, name string)
 		setupStorageFn func(m *mock.Mock, name string)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -214,6 +220,7 @@ func TestMode1_Delete(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -234,21 +241,22 @@ func TestMode1_Delete(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			us.AssertNotCalled(t, "Delete", context.Background(), tt.input, func(ctx context.Context, obj runtime.Object) error { return nil }, &metav1.DeleteOptions{})
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode1_DeleteCollection(t *testing.T) {
 	type testCase struct {
-		name           string
 		input          *metav1.DeleteOptions
 		setupLegacyFn  func(m *mock.Mock, input *metav1.DeleteOptions)
 		setupStorageFn func(m *mock.Mock, input *metav1.DeleteOptions)
+		name           string
 		wantErr        bool
 	}
 	tests :=
@@ -271,6 +279,7 @@ func TestMode1_DeleteCollection(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -291,22 +300,23 @@ func TestMode1_DeleteCollection(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			us.AssertNotCalled(t, "DeleteCollection", context.Background(), tt.input, func(ctx context.Context, obj runtime.Object) error { return nil }, &metav1.DeleteOptions{})
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode1_Update(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock, input string)
 		setupStorageFn func(m *mock.Mock, input string)
 		setupGetFn     func(m *mock.Mock, input string)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -341,6 +351,7 @@ func TestMode1_Update(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -361,14 +372,15 @@ func TestMode1_Update(t *testing.T) {
 
 			dw := NewDualWriter(Mode1, ls, us)
 
-		obj, _, err := dw.Update(context.Background(), tt.input, UpdatedObjInfoObj{}, func(ctx context.Context, obj runtime.Object) error { return nil }, func(ctx context.Context, obj, old runtime.Object) error { return nil }, false, &metav1.UpdateOptions{})
+			obj, _, err := dw.Update(context.Background(), tt.input, updatedObjInfoObj{}, func(ctx context.Context, obj runtime.Object) error { return nil }, func(ctx context.Context, obj, old runtime.Object) error { return nil }, false, &metav1.UpdateOptions{})
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }

pkg/apiserver/rest/dualwriter_mode2.go

@@ -18,8 +18,8 @@ import (
 type DualWriterMode2 struct {
 	Storage Storage
 	Legacy  LegacyStorage
-	Log     klog.Logger
 	*dualWriterMetrics
+	Log klog.Logger
 }
 
 // NewDualWriterMode2 returns a new DualWriter in mode 2.
@@ -36,40 +36,29 @@ func (d *DualWriterMode2) Mode() DualWriterMode {
 }
 
 // Create overrides the behavior of the generic DualWriter and writes to LegacyStorage and Storage.
-func (d *DualWriterMode2) Create(ctx context.Context, obj runtime.Object, createValidation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
+func (d *DualWriterMode2) Create(ctx context.Context, original runtime.Object, createValidation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
 	log := d.Log.WithValues("kind", options.Kind)
 	ctx = klog.NewContext(ctx, log)
 
-	created, err := d.Legacy.Create(ctx, obj, createValidation, options)
+	created, err := d.Legacy.Create(ctx, original, createValidation, options)
 	if err != nil {
 		log.Error(err, "unable to create object in legacy storage")
 		return created, err
 	}
 
-	accessorCreated, err := meta.Accessor(created)
+	createdLegacy, err := enrichLegacyObject(original, created, true)
 	if err != nil {
-		return created, err
+		return createdLegacy, err
 	}
 
-	accessorOld, err := meta.Accessor(obj)
-	if err != nil {
-		return created, err
-	}
-
-	enrichObject(accessorOld, accessorCreated)
-
-	// create method expects an empty resource version
-	accessorCreated.SetResourceVersion("")
-	accessorCreated.SetUID("")
-
 	rsp, err := d.Storage.Create(ctx, created, createValidation, options)
 	if err != nil {
-		log.WithValues("name", accessorCreated.GetName(), "resourceVersion", accessorCreated.GetResourceVersion()).Error(err, "unable to create object in storage")
-	}
+		log.WithValues("name").Error(err, "unable to create object in storage")
 		return rsp, err
+	}
+	return rsp, nil
 }
 
-// Get overrides the behavior of the generic DualWriter.
 // It retrieves an object from Storage if possible, and if not it falls back to LegacyStorage.
 func (d *DualWriterMode2) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
 	log := d.Log.WithValues("name", name, "resourceVersion", options.ResourceVersion, "kind", options.Kind)
@@ -199,7 +188,7 @@ func (d *DualWriterMode2) Update(ctx context.Context, name string, objInfo rest.
 	log := d.Log.WithValues("name", name, "kind", options.Kind)
 	ctx = klog.NewContext(ctx, log)
 
-	// get foundObj and new (updated) object so they can be stored in legacy store
+	// get foundObj and (updated) object so they can be stored in legacy store
 	foundObj, err := d.Storage.Get(ctx, name, &metav1.GetOptions{})
 	if err != nil {
 		if !apierrors.IsNotFound(err) {
@@ -222,30 +211,18 @@ func (d *DualWriterMode2) Update(ctx context.Context, name string, objInfo rest.
 		return obj, created, err
 	}
 
-	// if the object is found, create a new updateWrapper with the object found
 	if foundObj != nil {
-		accessorOld, err := meta.Accessor(foundObj)
+		obj, err = enrichLegacyObject(foundObj, obj, false)
 		if err != nil {
-			log.Error(err, "unable to get accessor for original updated object")
+			return obj, false, err
 		}
 
-		accessor, err := meta.Accessor(obj)
-		if err != nil {
-			log.Error(err, "unable to get accessor for updated object")
-		}
-
-		enrichObject(accessorOld, accessor)
-
-		accessor.SetResourceVersion(accessorOld.GetResourceVersion())
-		accessor.SetUID(accessorOld.GetUID())
-
 		objInfo = &updateWrapper{
 			upstream: objInfo,
 			updated:  obj,
 		}
 	}
-	// TODO: relies on GuaranteedUpdate creating the object if
-	// it doesn't exist: https://github.com/grafana/grafana/pull/85206
+
 	return d.Storage.Update(ctx, name, objInfo, createValidation, updateValidation, forceAllowCreate, options)
 }
 
@@ -306,15 +283,37 @@ func parseList(legacyList []runtime.Object) (metainternalversion.ListOptions, ma
 	return options, indexMap, nil
 }
 
-func enrichObject(accessorO, accessorC metav1.Object) {
-	accessorC.SetLabels(accessorO.GetLabels())
+func enrichLegacyObject(originalObj, returnedObj runtime.Object, created bool) (runtime.Object, error) {
+	accessorReturned, err := meta.Accessor(returnedObj)
+	if err != nil {
+		return nil, err
+	}
 
-	ac := accessorC.GetAnnotations()
+	accessorOriginal, err := meta.Accessor(originalObj)
+	if err != nil {
+		return nil, err
+	}
+
+	accessorReturned.SetLabels(accessorOriginal.GetLabels())
+
+	ac := accessorReturned.GetAnnotations()
 	if ac == nil {
 		ac = map[string]string{}
 	}
-	for k, v := range accessorO.GetAnnotations() {
+	for k, v := range accessorOriginal.GetAnnotations() {
 		ac[k] = v
 	}
-	accessorC.SetAnnotations(ac)
+	accessorReturned.SetAnnotations(ac)
+
+	// if the object is created, we need to reset the resource version and UID
+	// create method expects an empty resource version
+	if created {
+		accessorReturned.SetResourceVersion("")
+		accessorReturned.SetUID("")
+		return returnedObj, nil
+	}
+	// otherwise, we propagate the original RV and UID
+	accessorReturned.SetResourceVersion(accessorOriginal.GetResourceVersion())
+	accessorReturned.SetUID(accessorOriginal.GetUID())
+	return returnedObj, nil
 }

pkg/apiserver/rest/dualwriter_mode2_test.go

@@ -13,6 +13,8 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apiserver/pkg/apis/example"
 )
 
 var createFn = func(context.Context, runtime.Object) error { return nil }
@@ -21,10 +23,10 @@ var exampleOption = &metainternalversion.ListOptions{}
 
 func TestMode2_Create(t *testing.T) {
 	type testCase struct {
-		name           string
 		input          runtime.Object
 		setupLegacyFn  func(m *mock.Mock, input runtime.Object)
 		setupStorageFn func(m *mock.Mock, input runtime.Object)
+		name           string
 		wantErr        bool
 	}
 	tests :=
@@ -50,6 +52,7 @@ func TestMode2_Create(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -70,22 +73,23 @@ func TestMode2_Create(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			assert.Equal(t, exampleObj, obj)
 			accessor, err := meta.Accessor(obj)
 			assert.NoError(t, err)
 			assert.Equal(t, accessor.GetResourceVersion(), "")
+		})
 	}
 }
 
 func TestMode2_Get(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock, input string)
 		setupStorageFn func(m *mock.Mock, input string)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -124,6 +128,7 @@ func TestMode2_Get(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -144,20 +149,21 @@ func TestMode2_Get(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode2_List(t *testing.T) {
 	type testCase struct {
-		name           string
 		inputLegacy    *metainternalversion.ListOptions
 		setupLegacyFn  func(m *mock.Mock)
 		setupStorageFn func(m *mock.Mock)
+		name           string
 		wantErr        bool
 	}
 	tests :=
@@ -175,6 +181,7 @@ func TestMode2_List(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -195,19 +202,19 @@ func TestMode2_List(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
-
 			assert.Equal(t, exampleList, obj)
+		})
 	}
 }
 
 func TestMode2_Delete(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock, input string)
 		setupStorageFn func(m *mock.Mock, input string)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -267,6 +274,7 @@ func TestMode2_Delete(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -287,20 +295,21 @@ func TestMode2_Delete(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			assert.Equal(t, obj, exampleObj)
 			assert.NotEqual(t, obj, anotherObj)
+		})
 	}
 }
 
 func TestMode2_DeleteCollection(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
 		setupLegacyFn  func(m *mock.Mock)
 		setupStorageFn func(m *mock.Mock)
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -337,6 +346,7 @@ func TestMode2_DeleteCollection(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -357,21 +367,21 @@ func TestMode2_DeleteCollection(t *testing.T) {
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
-
 			assert.Equal(t, exampleList, obj)
+		})
 	}
 }
 
 func TestMode2_Update(t *testing.T) {
 	type testCase struct {
-		name           string
-		input          string
+		expectedObj    runtime.Object
 		setupLegacyFn  func(m *mock.Mock, input string)
 		setupStorageFn func(m *mock.Mock, input string)
 		setupGetFn     func(m *mock.Mock, input string)
-		expectedObj    runtime.Object
+		name           string
+		input          string
 		wantErr        bool
 	}
 	tests :=
@@ -440,6 +450,7 @@ func TestMode2_Update(t *testing.T) {
 		}
 
 	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
 			l := (LegacyStorage)(nil)
 			s := (Storage)(nil)
 			m := &mock.Mock{}
@@ -460,14 +471,215 @@ func TestMode2_Update(t *testing.T) {
 
 			dw := NewDualWriter(Mode2, ls, us)
 
-		obj, _, err := dw.Update(context.Background(), tt.input, UpdatedObjInfoObj{}, func(ctx context.Context, obj runtime.Object) error { return nil }, func(ctx context.Context, obj, old runtime.Object) error { return nil }, false, &metav1.UpdateOptions{})
+			obj, _, err := dw.Update(context.Background(), tt.input, updatedObjInfoObj{}, func(ctx context.Context, obj runtime.Object) error { return nil }, func(ctx context.Context, obj, old runtime.Object) error { return nil }, false, &metav1.UpdateOptions{})
 
 			if tt.wantErr {
 				assert.Error(t, err)
-			continue
+				return
 			}
 
 			assert.Equal(t, tt.expectedObj, obj)
 			assert.NotEqual(t, anotherObj, obj)
+		})
+	}
+}
+func TestEnrichReturnedObject(t *testing.T) {
+	testCase := []struct {
+		inputOriginal  runtime.Object
+		inputReturned  runtime.Object
+		expectedObject runtime.Object
+		name           string
+		isCreated      bool
+		wantErr        bool
+	}{
+		{
+			name: "create: original object does not have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			isCreated: true,
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "", UID: types.UID("")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "create: returned object does not have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			isCreated: true,
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "", UID: types.UID(""), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "create: both objects have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label2": "2"}, Annotations: map[string]string{"annotation2": "2"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			isCreated: true,
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "", UID: types.UID(""), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "create: both objects have labels and annotations with duplicated keys",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label1": "11"}, Annotations: map[string]string{"annotation1": "11"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			isCreated: true,
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "", UID: types.UID(""), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "update: original object does not have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "update: returned object does not have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6")},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "update: both objects have labels and annotations",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label2": "2"}, Annotations: map[string]string{"annotation2": "2"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name: "update: both objects have labels and annotations with duplicated keys",
+			inputOriginal: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			inputReturned: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "2", UID: types.UID("6"), Labels: map[string]string{"label1": "11"}, Annotations: map[string]string{"annotation1": "11"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+			expectedObject: &example.Pod{
+				TypeMeta:   metav1.TypeMeta{Kind: "foo"},
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", ResourceVersion: "1", UID: types.UID("5"), Labels: map[string]string{"label1": "1"}, Annotations: map[string]string{"annotation1": "1"}},
+				Spec:       example.PodSpec{}, Status: example.PodStatus{},
+			},
+		},
+		{
+			name:           "original object does not exist",
+			inputOriginal:  nil,
+			inputReturned:  &example.Pod{},
+			expectedObject: nil,
+			wantErr:        true,
+		},
+		{
+			name:           "returned object does not exist",
+			inputOriginal:  &example.Pod{},
+			inputReturned:  nil,
+			expectedObject: nil,
+			wantErr:        true,
+		},
+	}
+
+	for _, tt := range testCase {
+		t.Run(tt.name, func(t *testing.T) {
+			returned, err := enrichLegacyObject(tt.inputOriginal, tt.inputReturned, tt.isCreated)
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			}
+
+			accessorReturned, err := meta.Accessor(returned)
+			assert.NoError(t, err)
+
+			accessorExpected, err := meta.Accessor(tt.expectedObject)
+			assert.NoError(t, err)
+
+			assert.Equal(t, accessorExpected.GetLabels(), accessorReturned.GetLabels())
+
+			returnedAnnotations := accessorReturned.GetAnnotations()
+			expectedAnnotations := accessorExpected.GetAnnotations()
+			for k, v := range expectedAnnotations {
+				assert.Equal(t, v, returnedAnnotations[k])
+			}
+
+			assert.Equal(t, accessorExpected.GetResourceVersion(), accessorReturned.GetResourceVersion())
+			assert.Equal(t, accessorExpected.GetUID(), accessorReturned.GetUID())
+		})
 	}
 }

pkg/apiserver/rest/dualwriter_mode3.go

@@ -14,8 +14,8 @@ import (
 type DualWriterMode3 struct {
 	Legacy  LegacyStorage
 	Storage Storage
-	Log     klog.Logger
 	*dualWriterMetrics
+	Log klog.Logger
 }
 
 // NewDualWriterMode3 returns a new DualWriter in mode 3.

pkg/apiserver/rest/dualwriter_mode4.go

@@ -13,8 +13,8 @@ import (
 type DualWriterMode4 struct {
 	Legacy  LegacyStorage
 	Storage Storage
-	Log     klog.Logger
 	*dualWriterMetrics
+	Log klog.Logger
 }
 
 // NewDualWriterMode4 returns a new DualWriter in mode 4.

pkg/apiserver/rest/spy_client.go

@@ -1,234 +0,0 @@
-package rest
-
-import (
-	"context"
-
-	"k8s.io/apimachinery/pkg/api/meta"
-	metainternalversion "k8s.io/apimachinery/pkg/apis/meta/internalversion"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/runtime"
-	"k8s.io/apimachinery/pkg/runtime/schema"
-	"k8s.io/apiserver/pkg/apis/example"
-	"k8s.io/apiserver/pkg/registry/rest"
-	"k8s.io/klog/v2"
-)
-
-// Unified Storage Spy
-
-type StorageSpyClient interface {
-	Storage
-
-	//Counts returns the number of times a certain method was called
-	Counts(string) int
-}
-
-type StorageSpy struct {
-	counts map[string]int
-}
-
-type spyStorageClient struct {
-	Storage
-	spy *StorageSpy
-}
-
-func (s *StorageSpy) record(seen string) {
-	s.counts[seen]++
-}
-
-func NewStorageSpyClient(s Storage) StorageSpyClient {
-	return &spyStorageClient{s, &StorageSpy{
-		counts: map[string]int{},
-	}}
-}
-
-func (c *spyStorageClient) Counts(method string) int {
-	return c.spy.counts[method]
-}
-
-//nolint:golint,unused
-type spyStorageShim struct {
-	Storage
-	spy *StorageSpy
-}
-
-//nolint:golint,unused
-type spyLegacyStorageShim struct {
-	LegacyStorage
-	spy *StorageSpy
-}
-
-func (c *spyStorageClient) Create(ctx context.Context, obj runtime.Object, valitation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
-	c.spy.record("Storage.Create")
-	klog.Info("method: Storage.Create")
-	return &dummyObject{}, nil
-}
-
-func (c *spyStorageClient) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
-	c.spy.record("Storage.Get")
-	klog.Info("method: Storage.Get")
-	return &example.Pod{}, nil
-}
-
-func (c *spyStorageClient) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
-	c.spy.record("Storage.List")
-	klog.Info("method: Storage.List")
-
-	i1 := dummyObject{Foo: "Storage field 1"}
-	accessor, err := meta.Accessor(&i1)
-	if err != nil {
-		return nil, err
-	}
-	accessor.SetName("Item 1")
-
-	i2 := dummyObject{Foo: "Storage field 2"}
-	accessor, err = meta.Accessor(&i2)
-	if err != nil {
-		return nil, err
-	}
-	accessor.SetName("Item 2")
-
-	return &dummyList{Items: []dummyObject{i1, i2}}, nil
-}
-
-type UpdatedObjInfoObj struct{}
-
-func (u UpdatedObjInfoObj) UpdatedObject(ctx context.Context, oldObj runtime.Object) (newObj runtime.Object, err error) {
-	return &example.Pod{}, nil
-}
-
-func (u UpdatedObjInfoObj) Preconditions() *metav1.Preconditions { return &metav1.Preconditions{} }
-
-func (c *spyStorageClient) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool, options *metav1.UpdateOptions) (runtime.Object, bool, error) {
-	c.spy.record("Storage.Update")
-	klog.Info("method: Storage.Update")
-	return &example.Pod{}, false, nil
-}
-
-func (c *spyStorageClient) Delete(ctx context.Context, name string, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions) (runtime.Object, bool, error) {
-	c.spy.record("Storage.Delete")
-	klog.Info("method: Storage.Delete")
-	return nil, false, nil
-}
-
-func (c *spyStorageClient) DeleteCollection(ctx context.Context, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions, listOptions *metainternalversion.ListOptions) (runtime.Object, error) {
-	c.spy.record("Storage.DeleteCollection")
-	klog.Info("method: Storage.DeleteCollection")
-	return nil, nil
-}
-
-// LegacyStorage Spy
-
-type LegacyStorageSpyClient interface {
-	LegacyStorage
-
-	//Counts returns the number of times a certain method was called
-	Counts(string) int
-}
-
-type LegacyStorageSpy struct {
-	counts map[string]int //nolint:golint,unused
-}
-
-type spyLegacyStorageClient struct {
-	LegacyStorage
-	spy *StorageSpy
-}
-
-//nolint:golint,unused
-func (s *LegacyStorageSpy) record(seen string) {
-	s.counts[seen]++
-}
-
-func NewLegacyStorageSpyClient(ls LegacyStorage) LegacyStorageSpyClient {
-	return &spyLegacyStorageClient{ls, &StorageSpy{
-		counts: map[string]int{},
-	}}
-}
-
-func (c *spyLegacyStorageClient) Counts(method string) int {
-	return c.spy.counts[method]
-}
-
-func (c *spyLegacyStorageClient) Create(ctx context.Context, obj runtime.Object, valitation rest.ValidateObjectFunc, options *metav1.CreateOptions) (runtime.Object, error) {
-	c.spy.record("LegacyStorage.Create")
-	klog.Info("method: LegacyStorage.Create")
-	return &dummyObject{}, nil
-}
-
-func (c *spyLegacyStorageClient) Get(ctx context.Context, name string, options *metav1.GetOptions) (runtime.Object, error) {
-	c.spy.record("LegacyStorage.Get")
-	klog.Info("method: LegacyStorage.Get")
-	return &example.Pod{}, nil
-}
-
-func (c *spyLegacyStorageClient) NewList() runtime.Object {
-	// stub for now so that spyLegacyStorageClient implements rest.Lister
-	return nil
-}
-
-func (c *spyLegacyStorageClient) List(ctx context.Context, options *metainternalversion.ListOptions) (runtime.Object, error) {
-	c.spy.record("LegacyStorage.List")
-	klog.Info("method: LegacyStorage.List")
-
-	i1 := dummyObject{Foo: "Legacy field 1"}
-	accessor, err := meta.Accessor(&i1)
-	if err != nil {
-		return nil, err
-	}
-	accessor.SetName("Item 1")
-
-	i3 := dummyObject{Foo: "Legacy field 3"}
-	accessor, err = meta.Accessor(&i3)
-	if err != nil {
-		return nil, err
-	}
-	accessor.SetName("Item 3")
-
-	return &dummyList{Items: []dummyObject{i1, i3}}, nil
-}
-
-func (c *spyLegacyStorageClient) Update(ctx context.Context, name string, objInfo rest.UpdatedObjectInfo, createValidation rest.ValidateObjectFunc, updateValidation rest.ValidateObjectUpdateFunc, forceAllowCreate bool, options *metav1.UpdateOptions) (runtime.Object, bool, error) {
-	c.spy.record("LegacyStorage.Update")
-	klog.Info("method: LegacyStorage.Update")
-	return &example.Pod{}, false, nil
-}
-
-func (c *spyLegacyStorageClient) Delete(ctx context.Context, name string, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions) (runtime.Object, bool, error) {
-	c.spy.record("LegacyStorage.Delete")
-	klog.Info("method: LegacyStorage.Delete")
-	return nil, false, nil
-}
-
-func (c *spyLegacyStorageClient) DeleteCollection(ctx context.Context, deleteValidation rest.ValidateObjectFunc, options *metav1.DeleteOptions, listOptions *metainternalversion.ListOptions) (runtime.Object, error) {
-	c.spy.record("LegacyStorage.DeleteCollection")
-	klog.Info("method: LegacyStorage.DeleteCollection")
-	return nil, nil
-}
-
-type dummyList struct {
-	metav1.TypeMeta `json:",inline"`
-	metav1.ListMeta `json:"metadata,omitempty"`
-	Items           []dummyObject `json:"items,omitempty"`
-}
-
-type dummyObject struct {
-	metav1.TypeMeta   `json:",inline"`
-	Foo               string
-	metav1.ObjectMeta `json:"metadata,omitempty"`
-}
-
-func (d *dummyList) GetObjectKind() schema.ObjectKind {
-	return nil
-}
-
-func (d *dummyList) DeepCopyObject() runtime.Object {
-	return nil
-}
-
-func (d *dummyObject) GetObjectKind() schema.ObjectKind {
-	return nil
-}
-
-func (d *dummyObject) DeepCopyObject() runtime.Object {
-	return nil
-}

pkg/apiserver/rest/storage_mocks_test.go

@@ -141,3 +141,12 @@ func (m storageMock) DeleteCollection(ctx context.Context, deleteValidation rest
 	}
 	return args.Get(0).(runtime.Object), args.Error(1)
 }
+
+type updatedObjInfoObj struct{}
+
+func (u updatedObjInfoObj) UpdatedObject(ctx context.Context, oldObj runtime.Object) (newObj runtime.Object, err error) { // nolint:staticcheck
+	// nolint:staticcheck
+	oldObj = exampleObj
+	return oldObj, nil
+}
+func (u updatedObjInfoObj) Preconditions() *metav1.Preconditions { return &metav1.Preconditions{} }