pod: Add test for exit code

pkg/controller.v2/controller_pod_test.go

@@ -236,3 +236,78 @@ func TestRestartPolicy(t *testing.T) {
 		}
 	}
 }
+
+func TestExitCode(t *testing.T) {
+	// Prepare the clientset and controller for the test.
+	kubeClientSet := kubeclientset.NewForConfigOrDie(&rest.Config{
+		Host: "",
+		ContentConfig: rest.ContentConfig{
+			GroupVersion: &v1.SchemeGroupVersion,
+		},
+	},
+	)
+	config := &rest.Config{
+		Host: "",
+		ContentConfig: rest.ContentConfig{
+			GroupVersion: &tfv1alpha2.SchemeGroupVersion,
+		},
+	}
+	tfJobClientSet := tfjobclientset.NewForConfigOrDie(config)
+	ctr, kubeInformerFactory, _ := newTFJobController(config, kubeClientSet, tfJobClientSet, controller.NoResyncPeriodFunc)
+	fakePodControl := &controller.FakePodControl{}
+	ctr.podControl = fakePodControl
+	ctr.tfJobInformerSynced = alwaysReady
+	ctr.podInformerSynced = alwaysReady
+	ctr.serviceInformerSynced = alwaysReady
+	tfJobIndexer := ctr.tfJobInformer.GetIndexer()
+	podIndexer := kubeInformerFactory.Core().V1().Pods().Informer().GetIndexer()
+
+	stopCh := make(chan struct{})
+	run := func(<-chan struct{}) {
+		ctr.Run(threadCount, stopCh)
+	}
+	go run(stopCh)
+
+	ctr.updateStatusHandler = func(tfJob *tfv1alpha2.TFJob) error {
+		return nil
+	}
+
+	tfJob := newTFJob(1, 0)
+	tfJob.Spec.TFReplicaSpecs[tfv1alpha2.TFReplicaTypeWorker].RestartPolicy = tfv1alpha2.RestartPolicyExitCode
+	unstructured, err := convertTFJobToUnstructured(tfJob)
+	if err != nil {
+		t.Errorf("Failed to convert the TFJob to Unstructured: %v", err)
+	}
+
+	if err := tfJobIndexer.Add(unstructured); err != nil {
+		t.Errorf("Failed to add tfjob to tfJobIndexer: %v", err)
+	}
+	pod := newPod(tfJob, labelWorker, 0, t)
+	pod.Status.Phase = v1.PodFailed
+	pod.Spec.Containers = append(pod.Spec.Containers, v1.Container{})
+	pod.Status.ContainerStatuses = append(pod.Status.ContainerStatuses, v1.ContainerStatus{
+		Name: tfv1alpha2.DefaultContainerName,
+		State: v1.ContainerState{
+			Terminated: &v1.ContainerStateTerminated{
+				ExitCode: 130,
+			},
+		},
+	})
+
+	podIndexer.Add(pod)
+	_, err = ctr.syncTFJob(getKey(tfJob, t))
+	if err != nil {
+		t.Errorf("%s: unexpected error when syncing jobs %v", tfJob.Name, err)
+	}
+
+	found := false
+	for _, deletedPodName := range fakePodControl.DeletePodName {
+		if deletedPodName == pod.Name {
+			found = true
+		}
+	}
+	if !found {
+		t.Errorf("Failed to delete pod %s", pod.Name)
+	}
+	close(stopCh)
+}

pkg/controller.v2/pod_control.go

@@ -21,7 +21,6 @@ package controller
 
 import (
 	"fmt"
-	"sync"
 
 	"github.com/golang/glog"
 	"k8s.io/api/core/v1"
@@ -164,91 +163,3 @@ func (r RealPodControl) DeletePod(namespace string, podID string, object runtime
 	}
 	return nil
 }
-
-type FakePodControl struct {
-	sync.Mutex
-	Templates       []v1.PodTemplateSpec
-	ControllerRefs  []metav1.OwnerReference
-	DeletePodName   []string
-	Patches         [][]byte
-	Err             error
-	CreateLimit     int
-	CreateCallCount int
-}
-
-var _ controller.PodControlInterface = &FakePodControl{}
-
-func (f *FakePodControl) PatchPod(namespace, name string, data []byte) error {
-	f.Lock()
-	defer f.Unlock()
-	f.Patches = append(f.Patches, data)
-	if f.Err != nil {
-		return f.Err
-	}
-	return nil
-}
-
-func (f *FakePodControl) CreatePods(namespace string, spec *v1.PodTemplateSpec, object runtime.Object) error {
-	f.Lock()
-	defer f.Unlock()
-	f.CreateCallCount++
-	if f.CreateLimit != 0 && f.CreateCallCount > f.CreateLimit {
-		return fmt.Errorf("Not creating pod, limit %d already reached (create call %d)", f.CreateLimit, f.CreateCallCount)
-	}
-	f.Templates = append(f.Templates, *spec)
-	if f.Err != nil {
-		return f.Err
-	}
-	return nil
-}
-
-func (f *FakePodControl) CreatePodsWithControllerRef(namespace string, spec *v1.PodTemplateSpec, object runtime.Object, controllerRef *metav1.OwnerReference) error {
-	f.Lock()
-	defer f.Unlock()
-	f.CreateCallCount++
-	if f.CreateLimit != 0 && f.CreateCallCount > f.CreateLimit {
-		return fmt.Errorf("Not creating pod, limit %d already reached (create call %d)", f.CreateLimit, f.CreateCallCount)
-	}
-	f.Templates = append(f.Templates, *spec)
-	f.ControllerRefs = append(f.ControllerRefs, *controllerRef)
-	if f.Err != nil {
-		return f.Err
-	}
-	return nil
-}
-
-func (f *FakePodControl) CreatePodsOnNode(nodeName, namespace string, template *v1.PodTemplateSpec, object runtime.Object, controllerRef *metav1.OwnerReference) error {
-	f.Lock()
-	defer f.Unlock()
-	f.CreateCallCount++
-	if f.CreateLimit != 0 && f.CreateCallCount > f.CreateLimit {
-		return fmt.Errorf("Not creating pod, limit %d already reached (create call %d)", f.CreateLimit, f.CreateCallCount)
-	}
-	f.Templates = append(f.Templates, *template)
-	f.ControllerRefs = append(f.ControllerRefs, *controllerRef)
-	if f.Err != nil {
-		return f.Err
-	}
-	return nil
-}
-
-func (f *FakePodControl) DeletePod(namespace string, podID string, object runtime.Object) error {
-	f.Lock()
-	defer f.Unlock()
-	f.DeletePodName = append(f.DeletePodName, podID)
-	if f.Err != nil {
-		return f.Err
-	}
-	return nil
-}
-
-func (f *FakePodControl) Clear() {
-	f.Lock()
-	defer f.Unlock()
-	f.DeletePodName = []string{}
-	f.Templates = []v1.PodTemplateSpec{}
-	f.ControllerRefs = []metav1.OwnerReference{}
-	f.Patches = [][]byte{}
-	f.CreateLimit = 0
-	f.CreateCallCount = 0
-}