From 90367729a3b56bab1f95ad4c8c5569826d779053 Mon Sep 17 00:00:00 2001
From: marosset <marosset@microsoft.com>
Date: Fri, 17 Apr 2020 00:56:36 +0000
Subject: [PATCH] Adding/updating kubelet/kuberuntime tests

---
 pkg/kubelet/container/testing/fake_runtime.go |  9 +++-
 pkg/kubelet/images/image_manager_test.go      | 47 ++++++++++++++++++-
 .../kuberuntime/kuberuntime_image_test.go     | 24 +++++++++-
 3 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/pkg/kubelet/container/testing/fake_runtime.go b/pkg/kubelet/container/testing/fake_runtime.go
index 1824b8b84b7..87171752f1d 100644
--- a/pkg/kubelet/container/testing/fake_runtime.go
+++ b/pkg/kubelet/container/testing/fake_runtime.go
@@ -25,7 +25,7 @@ import (
 	"sync"
 	"time"
 
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/util/flowcontrol"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
@@ -300,6 +300,13 @@ func (f *FakeRuntime) PullImage(image kubecontainer.ImageSpec, pullSecrets []v1.
 	defer f.Unlock()
 
 	f.CalledFunctions = append(f.CalledFunctions, "PullImage")
+	if f.Err == nil {
+		i := kubecontainer.Image{
+			ID:   image.Image,
+			Spec: image,
+		}
+		f.ImageList = append(f.ImageList, i)
+	}
 	return image.Image, f.Err
 }
 
diff --git a/pkg/kubelet/images/image_manager_test.go b/pkg/kubelet/images/image_manager_test.go
index d2e27dcd73a..dff58f6c42e 100644
--- a/pkg/kubelet/images/image_manager_test.go
+++ b/pkg/kubelet/images/image_manager_test.go
@@ -22,7 +22,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/clock"
 	"k8s.io/client-go/tools/record"
@@ -202,3 +202,48 @@ func TestApplyDefaultImageTag(t *testing.T) {
 		}
 	}
 }
+
+func TestPullAndListImageWithPodAnnotations(t *testing.T) {
+	pod := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            "test_pod",
+			Namespace:       "test-ns",
+			UID:             "bar",
+			ResourceVersion: "42",
+			SelfLink:        "/api/v1/pods/foo",
+			Annotations: map[string]string{
+				"kubernetes.io/runtimehandler": "handler_name",
+			},
+		}}
+	c := pullerTestCase{ // pull missing image
+		containerImage: "missing_image",
+		policy:         v1.PullIfNotPresent,
+		inspectErr:     nil,
+		pullerErr:      nil,
+		expected: []pullerExpects{
+			{[]string{"GetImageRef", "PullImage"}, nil},
+		}}
+
+	useSerializedEnv := true
+	puller, fakeClock, fakeRuntime, container := pullerTestEnv(c, useSerializedEnv)
+	fakeRuntime.CalledFunctions = nil
+	fakeRuntime.ImageList = []Image{}
+	fakeClock.Step(time.Second)
+
+	_, _, err := puller.EnsureImageExists(pod, container, nil, nil)
+	assert.NoError(t, fakeRuntime.AssertCalls(c.expected[0].calls), "tick=%d", 0)
+	assert.Equal(t, c.expected[0].err, err, "tick=%d", 0)
+
+	images, _ := fakeRuntime.ListImages()
+	assert.Equal(t, 1, len(images), "ListImages() count")
+
+	image := images[0]
+	assert.Equal(t, "missing_image:latest", image.ID, "Image ID")
+
+	expectedAnnotations := []Annotation{
+		{
+			Name:  "kubernetes.io/runtimehandler",
+			Value: "handler_name",
+		}}
+	assert.Equal(t, expectedAnnotations, image.Spec.Annotations, "image spec annotations")
+}
diff --git a/pkg/kubelet/kuberuntime/kuberuntime_image_test.go b/pkg/kubelet/kuberuntime/kuberuntime_image_test.go
index bb210c5bb22..9ef2d571db6 100644
--- a/pkg/kubelet/kuberuntime/kuberuntime_image_test.go
+++ b/pkg/kubelet/kuberuntime/kuberuntime_image_test.go
@@ -23,7 +23,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"k8s.io/api/core/v1"
+	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
 	"k8s.io/kubernetes/pkg/credentialprovider"
@@ -168,6 +168,26 @@ func TestPullWithSecrets(t *testing.T) {
 
 		_, err = fakeManager.PullImage(kubecontainer.ImageSpec{Image: test.imageName}, test.passedSecrets, nil)
 		require.NoError(t, err)
-		fakeImageService.AssertImagePulledWithAuth(t, &runtimeapi.ImageSpec{Image: test.imageName}, test.expectedAuth, description)
+		fakeImageService.AssertImagePulledWithAuth(t, &runtimeapi.ImageSpec{Image: test.imageName, Annotations: make(map[string]string)}, test.expectedAuth, description)
 	}
 }
+
+func TestPullThenListWithAnnotations(t *testing.T) {
+	_, _, fakeManager, err := createTestRuntimeManager()
+	assert.NoError(t, err)
+
+	imageSpec := kubecontainer.ImageSpec{
+		Image: "12345",
+		Annotations: []kubecontainer.Annotation{
+			{Name: "kubernetes.io/runtimehandler", Value: "handler_name"},
+		},
+	}
+
+	_, err = fakeManager.PullImage(imageSpec, nil, nil)
+	assert.NoError(t, err)
+
+	images, err := fakeManager.ListImages()
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(images))
+	assert.Equal(t, images[0].Spec, imageSpec)
+}