github.com/kubeflow/training-operator@v1.7.0/pkg/apis/kubeflow.org/v1/mpi_defaults_test.go (about)

     1  package v1
     2  
     3  import (
     4  	"reflect"
     5  	"testing"
     6  
     7  	corev1 "k8s.io/api/core/v1"
     8  	"k8s.io/utils/pointer"
     9  )
    10  
    11  func expectedMPIJob(cleanPodPolicy CleanPodPolicy, restartPolicy RestartPolicy) *MPIJob {
    12  	return &MPIJob{
    13  		Spec: MPIJobSpec{
    14  			CleanPodPolicy: &cleanPodPolicy,
    15  			RunPolicy: RunPolicy{
    16  				CleanPodPolicy: &cleanPodPolicy,
    17  			},
    18  			MPIReplicaSpecs: map[ReplicaType]*ReplicaSpec{
    19  				MPIJobReplicaTypeLauncher: {
    20  					Replicas:      pointer.Int32(1),
    21  					RestartPolicy: restartPolicy,
    22  					Template: corev1.PodTemplateSpec{
    23  						Spec: corev1.PodSpec{
    24  							Containers: []corev1.Container{
    25  								{
    26  									Name:  MPIJobDefaultContainerName,
    27  									Image: testImage,
    28  								},
    29  							},
    30  						},
    31  					},
    32  				},
    33  				MPIJobReplicaTypeWorker: {
    34  					Replicas:      pointer.Int32(0),
    35  					RestartPolicy: restartPolicy,
    36  					Template: corev1.PodTemplateSpec{
    37  						Spec: corev1.PodSpec{
    38  							Containers: []corev1.Container{
    39  								{
    40  									Name:  MPIJobDefaultContainerName,
    41  									Image: testImage,
    42  								},
    43  							},
    44  						},
    45  					},
    46  				},
    47  			},
    48  		},
    49  	}
    50  }
    51  
    52  func TestSetDefaults_MPIJob(t *testing.T) {
    53  	customRestartPolicy := RestartPolicyAlways
    54  
    55  	testCases := map[string]struct {
    56  		original *MPIJob
    57  		expected *MPIJob
    58  	}{
    59  		"set default replicas": {
    60  			original: &MPIJob{
    61  				Spec: MPIJobSpec{
    62  					CleanPodPolicy: CleanPodPolicyPointer(CleanPodPolicyRunning),
    63  					RunPolicy: RunPolicy{
    64  						CleanPodPolicy: CleanPodPolicyPointer(CleanPodPolicyRunning),
    65  					},
    66  					MPIReplicaSpecs: map[ReplicaType]*ReplicaSpec{
    67  						MPIJobReplicaTypeLauncher: {
    68  							RestartPolicy: customRestartPolicy,
    69  							Template: corev1.PodTemplateSpec{
    70  								Spec: corev1.PodSpec{
    71  									Containers: []corev1.Container{
    72  										{
    73  											Name:  MPIJobDefaultContainerName,
    74  											Image: testImage,
    75  										},
    76  									},
    77  								},
    78  							},
    79  						},
    80  						MPIJobReplicaTypeWorker: {
    81  							RestartPolicy: customRestartPolicy,
    82  							Template: corev1.PodTemplateSpec{
    83  								Spec: corev1.PodSpec{
    84  									Containers: []corev1.Container{
    85  										{
    86  											Name:  MPIJobDefaultContainerName,
    87  											Image: testImage,
    88  										},
    89  									},
    90  								},
    91  							},
    92  						},
    93  					},
    94  				},
    95  			},
    96  			expected: expectedMPIJob(CleanPodPolicyRunning, customRestartPolicy),
    97  		},
    98  		"set default clean pod policy": {
    99  			original: &MPIJob{
   100  				Spec: MPIJobSpec{
   101  					MPIReplicaSpecs: map[ReplicaType]*ReplicaSpec{
   102  						MPIJobReplicaTypeLauncher: {
   103  							Template: corev1.PodTemplateSpec{
   104  								Spec: corev1.PodSpec{
   105  									Containers: []corev1.Container{
   106  										{
   107  											Name:  MPIJobDefaultContainerName,
   108  											Image: testImage,
   109  										},
   110  									},
   111  								},
   112  							},
   113  						},
   114  						MPIJobReplicaTypeWorker: {
   115  							Template: corev1.PodTemplateSpec{
   116  								Spec: corev1.PodSpec{
   117  									Containers: []corev1.Container{
   118  										{
   119  											Name:  MPIJobDefaultContainerName,
   120  											Image: testImage,
   121  										},
   122  									},
   123  								},
   124  							},
   125  						},
   126  					},
   127  				},
   128  			},
   129  			expected: expectedMPIJob(CleanPodPolicyNone, MPIJobDefaultRestartPolicy),
   130  		},
   131  		"set default restart policy": {
   132  			original: &MPIJob{
   133  				Spec: MPIJobSpec{
   134  					MPIReplicaSpecs: map[ReplicaType]*ReplicaSpec{
   135  						MPIJobReplicaTypeLauncher: {
   136  							Template: corev1.PodTemplateSpec{
   137  								Spec: corev1.PodSpec{
   138  									Containers: []corev1.Container{
   139  										{
   140  											Name:  MPIJobDefaultContainerName,
   141  											Image: testImage,
   142  										},
   143  									},
   144  								},
   145  							},
   146  						},
   147  						MPIJobReplicaTypeWorker: {
   148  							Template: corev1.PodTemplateSpec{
   149  								Spec: corev1.PodSpec{
   150  									Containers: []corev1.Container{
   151  										{
   152  											Name:  MPIJobDefaultContainerName,
   153  											Image: testImage,
   154  										},
   155  									},
   156  								},
   157  							},
   158  						},
   159  					},
   160  				},
   161  			},
   162  			expected: expectedMPIJob(CleanPodPolicyNone, MPIJobDefaultRestartPolicy),
   163  		},
   164  	}
   165  	for name, tc := range testCases {
   166  		SetDefaults_MPIJob(tc.original)
   167  		if !reflect.DeepEqual(tc.original, tc.expected) {
   168  			t.Errorf("%s: Want\n%v; Got\n %v", name, tc.expected, tc.original)
   169  		}
   170  	}
   171  }