github.com/kubeflow/training-operator@v1.7.0/pkg/controller.v1/common/status_test.go (about)

     1  package common
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	apiv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
     8  	"github.com/stretchr/testify/assert"
     9  	corev1 "k8s.io/api/core/v1"
    10  	metaV1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    11  )
    12  
    13  func TestUpdateJobReplicaStatuses(t *testing.T) {
    14  	jobStatus := apiv1.JobStatus{}
    15  	initializeReplicaStatuses(&jobStatus, "worker")
    16  	_, ok := jobStatus.ReplicaStatuses["worker"]
    17  	// assert ReplicaStatus for "worker" exists
    18  	assert.True(t, ok)
    19  	setStatusForTest(&jobStatus, "worker", 2, 3, 1, 1)
    20  	// terminating pod should count as failed.
    21  	assert.Equal(t, jobStatus.ReplicaStatuses["worker"].Failed, int32(3))
    22  	assert.Equal(t, jobStatus.ReplicaStatuses["worker"].Succeeded, int32(3))
    23  	assert.Equal(t, jobStatus.ReplicaStatuses["worker"].Active, int32(1))
    24  }
    25  
    26  func setStatusForTest(jobStatus *apiv1.JobStatus, rtype apiv1.ReplicaType, failed, succeeded, active, terminating int32) {
    27  	pod := corev1.Pod{
    28  		Status: corev1.PodStatus{},
    29  	}
    30  	var i int32
    31  	for i = 0; i < failed; i++ {
    32  		pod.Status.Phase = corev1.PodFailed
    33  		updateJobReplicaStatuses(jobStatus, rtype, &pod)
    34  	}
    35  	for i = 0; i < succeeded; i++ {
    36  		pod.Status.Phase = corev1.PodSucceeded
    37  		updateJobReplicaStatuses(jobStatus, rtype, &pod)
    38  	}
    39  	for i = 0; i < active; i++ {
    40  		pod.Status.Phase = corev1.PodRunning
    41  		updateJobReplicaStatuses(jobStatus, rtype, &pod)
    42  	}
    43  	for i = 0; i < terminating; i++ {
    44  		pod.Status.Phase = corev1.PodRunning
    45  		deletionTimestamp := metaV1.NewTime(time.Now())
    46  		pod.DeletionTimestamp = &deletionTimestamp
    47  		updateJobReplicaStatuses(jobStatus, rtype, &pod)
    48  	}
    49  }