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

     1  package common
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/kubeflow/training-operator/pkg/core"
     7  
     8  	apiv1 "github.com/kubeflow/training-operator/pkg/apis/kubeflow.org/v1"
     9  	"github.com/stretchr/testify/assert"
    10  	corev1 "k8s.io/api/core/v1"
    11  	v1 "k8s.io/api/core/v1"
    12  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    13  )
    14  
    15  func TestCalculateServiceSliceSize(t *testing.T) {
    16  	type testCase struct {
    17  		services     []*corev1.Service
    18  		replicas     int
    19  		expectedSize int
    20  	}
    21  
    22  	services := []*corev1.Service{
    23  		{
    24  			ObjectMeta: metav1.ObjectMeta{
    25  				Labels: map[string]string{apiv1.ReplicaIndexLabel: "0"},
    26  			},
    27  		},
    28  		{
    29  			ObjectMeta: metav1.ObjectMeta{
    30  				Labels: map[string]string{apiv1.ReplicaIndexLabel: "1"},
    31  			},
    32  		},
    33  		{
    34  			ObjectMeta: metav1.ObjectMeta{
    35  				Labels: map[string]string{apiv1.ReplicaIndexLabel: "2"},
    36  			},
    37  		},
    38  	}
    39  
    40  	var testCases = []testCase{
    41  		{
    42  			services:     services,
    43  			replicas:     3,
    44  			expectedSize: 3,
    45  		},
    46  		{
    47  			services:     services,
    48  			replicas:     4,
    49  			expectedSize: 4,
    50  		},
    51  		{
    52  			services:     services,
    53  			replicas:     2,
    54  			expectedSize: 3,
    55  		},
    56  		{
    57  			services: append(services, &corev1.Service{
    58  				ObjectMeta: metav1.ObjectMeta{
    59  					Labels: map[string]string{apiv1.ReplicaIndexLabel: "4"},
    60  				},
    61  			}),
    62  			replicas:     3,
    63  			expectedSize: 5,
    64  		},
    65  	}
    66  
    67  	for _, tc := range testCases {
    68  		result := core.CalculateServiceSliceSize(tc.services, tc.replicas)
    69  		assert.Equal(t, tc.expectedSize, result)
    70  	}
    71  }
    72  
    73  func TestFilterServicesForReplicaType(t *testing.T) {
    74  	services := []*v1.Service{
    75  		{
    76  			ObjectMeta: metav1.ObjectMeta{
    77  				Name:   "a",
    78  				Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
    79  			},
    80  		},
    81  		{
    82  			ObjectMeta: metav1.ObjectMeta{
    83  				Name:   "b",
    84  				Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
    85  			},
    86  		},
    87  		{
    88  			ObjectMeta: metav1.ObjectMeta{
    89  				Name:   "c",
    90  				Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
    91  			},
    92  		},
    93  		{
    94  			ObjectMeta: metav1.ObjectMeta{
    95  				Name:   "d",
    96  				Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
    97  			},
    98  		},
    99  		{
   100  			ObjectMeta: metav1.ObjectMeta{
   101  				Name: "e",
   102  				Labels: map[string]string{
   103  					apiv1.ReplicaTypeLabel: "foo",
   104  				},
   105  			},
   106  		},
   107  	}
   108  	c := &JobController{}
   109  	got, err := c.FilterServicesForReplicaType(services, "foo")
   110  	if err != nil {
   111  		t.Fatalf("FilterPodsForReplicaType returned error: %v", err)
   112  	}
   113  	want := []*v1.Service{services[0], services[2], services[4]}
   114  	assert.Equal(t, want, got)
   115  }