From 33280abf426bfd90ccaa5375cff437ffd1fabf05 Mon Sep 17 00:00:00 2001
From: Random-Liu <taotaotheripper@gmail.com>
Date: Mon, 15 Feb 2016 13:12:36 -0800
Subject: [PATCH] Cleanup and add unit test for ShouldContainerBeRestarted

---
 pkg/kubelet/container/helpers.go      | 38 +++++++-------
 pkg/kubelet/container/helpers_test.go | 71 +++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 19 deletions(-)

diff --git a/pkg/kubelet/container/helpers.go b/pkg/kubelet/container/helpers.go
index f5346950b61..01baabcf1c8 100644
--- a/pkg/kubelet/container/helpers.go
+++ b/pkg/kubelet/container/helpers.go
@@ -46,28 +46,28 @@ type RuntimeHelper interface {
 // ShouldContainerBeRestarted checks whether a container needs to be restarted.
 // TODO(yifan): Think about how to refactor this.
 func ShouldContainerBeRestarted(container *api.Container, pod *api.Pod, podStatus *PodStatus) bool {
-	// Get all dead container status.
-	var resultStatus []*ContainerStatus
-	for _, containerStatus := range podStatus.ContainerStatuses {
-		if containerStatus.Name == container.Name && containerStatus.State == ContainerStateExited {
-			resultStatus = append(resultStatus, containerStatus)
-		}
+	// Get latest container status.
+	status := podStatus.FindContainerStatusByName(container.Name)
+	// If the container was never started before, we should start it.
+	// NOTE(random-liu): If all historical containers were GC'd, we'll also return true here.
+	if status == nil {
+		return true
 	}
-
-	// Check RestartPolicy for dead container.
-	if len(resultStatus) > 0 {
-		if pod.Spec.RestartPolicy == api.RestartPolicyNever {
-			glog.V(4).Infof("Already ran container %q of pod %q, do nothing", container.Name, format.Pod(pod))
+	// Check whether container is running
+	if status.State == ContainerStateRunning {
+		return false
+	}
+	// Check RestartPolicy for dead container
+	if pod.Spec.RestartPolicy == api.RestartPolicyNever {
+		glog.V(4).Infof("Already ran container %q of pod %q, do nothing", container.Name, format.Pod(pod))
+		return false
+	}
+	if pod.Spec.RestartPolicy == api.RestartPolicyOnFailure {
+		// Check the exit code.
+		if status.ExitCode == 0 {
+			glog.V(4).Infof("Already successfully ran container %q of pod %q, do nothing", container.Name, format.Pod(pod))
 			return false
 		}
-		if pod.Spec.RestartPolicy == api.RestartPolicyOnFailure {
-			// Check the exit code of last run. Note: This assumes the result is sorted
-			// by the created time in reverse order.
-			if resultStatus[0].ExitCode == 0 {
-				glog.V(4).Infof("Already successfully ran container %q of pod %q, do nothing", container.Name, format.Pod(pod))
-				return false
-			}
-		}
 	}
 	return true
 }
diff --git a/pkg/kubelet/container/helpers_test.go b/pkg/kubelet/container/helpers_test.go
index 600052722bb..81e520a4292 100644
--- a/pkg/kubelet/container/helpers_test.go
+++ b/pkg/kubelet/container/helpers_test.go
@@ -134,3 +134,74 @@ func TestExpandCommandAndArgs(t *testing.T) {
 
 	}
 }
+
+func TestShouldContainerBeRestarted(t *testing.T) {
+	pod := &api.Pod{
+		ObjectMeta: api.ObjectMeta{
+			UID:       "12345678",
+			Name:      "foo",
+			Namespace: "new",
+		},
+		Spec: api.PodSpec{
+			Containers: []api.Container{
+				{Name: "no-history"},
+				{Name: "alive"},
+				{Name: "succeed"},
+				{Name: "failed"},
+			},
+		},
+	}
+	podStatus := &PodStatus{
+		ID:        pod.UID,
+		Name:      pod.Name,
+		Namespace: pod.Namespace,
+		ContainerStatuses: []*ContainerStatus{
+			{
+				Name:  "alive",
+				State: ContainerStateRunning,
+			},
+			{
+				Name:     "succeed",
+				State:    ContainerStateExited,
+				ExitCode: 0,
+			},
+			{
+				Name:     "failed",
+				State:    ContainerStateExited,
+				ExitCode: 1,
+			},
+			{
+				Name:     "alive",
+				State:    ContainerStateExited,
+				ExitCode: 2,
+			},
+			{
+				Name:     "failed",
+				State:    ContainerStateExited,
+				ExitCode: 3,
+			},
+		},
+	}
+	policies := []api.RestartPolicy{
+		api.RestartPolicyNever,
+		api.RestartPolicyOnFailure,
+		api.RestartPolicyAlways,
+	}
+	expected := map[string][]bool{
+		"no-history": {true, true, true},
+		"alive":      {false, false, false},
+		"succeed":    {false, false, true},
+		"failed":     {false, true, true},
+	}
+	for _, c := range pod.Spec.Containers {
+		for i, policy := range policies {
+			pod.Spec.RestartPolicy = policy
+			e := expected[c.Name][i]
+			r := ShouldContainerBeRestarted(&c, pod, podStatus)
+			if r != e {
+				t.Errorf("Restart for container %q with restart policy %q expected %t, got %t",
+					c.Name, policy, e, r)
+			}
+		}
+	}
+}