sigs.k8s.io/kueue@v0.6.2/pkg/controller/admissionchecks/provisioning/controller_test.go (about)

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package provisioning
    18  
    19  import (
    20  	"testing"
    21  
    22  	"github.com/google/go-cmp/cmp"
    23  	"github.com/google/go-cmp/cmp/cmpopts"
    24  	corev1 "k8s.io/api/core/v1"
    25  	apierrors "k8s.io/apimachinery/pkg/api/errors"
    26  	apimeta "k8s.io/apimachinery/pkg/api/meta"
    27  	"k8s.io/apimachinery/pkg/api/resource"
    28  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    29  	"k8s.io/apimachinery/pkg/types"
    30  	autoscaling "k8s.io/autoscaler/cluster-autoscaler/apis/provisioningrequest/autoscaling.x-k8s.io/v1beta1"
    31  	"k8s.io/utils/ptr"
    32  	"sigs.k8s.io/controller-runtime/pkg/client"
    33  	"sigs.k8s.io/controller-runtime/pkg/reconcile"
    34  
    35  	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
    36  	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
    37  	"sigs.k8s.io/kueue/pkg/workload"
    38  )
    39  
    40  var (
    41  	wlCmpOptions = []cmp.Option{
    42  		cmpopts.EquateEmpty(),
    43  		cmpopts.IgnoreTypes(metav1.ObjectMeta{}, metav1.TypeMeta{}),
    44  		cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
    45  		cmpopts.IgnoreFields(kueue.AdmissionCheckState{}, "LastTransitionTime"),
    46  	}
    47  
    48  	reqCmpOptions = []cmp.Option{
    49  		cmpopts.EquateEmpty(),
    50  		cmpopts.IgnoreTypes(metav1.ObjectMeta{}, metav1.TypeMeta{}),
    51  		cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
    52  	}
    53  
    54  	tmplCmpOptions = []cmp.Option{
    55  		cmpopts.EquateEmpty(),
    56  		cmpopts.IgnoreTypes(metav1.ObjectMeta{}, metav1.TypeMeta{}),
    57  		cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
    58  		cmpopts.IgnoreFields(corev1.PodSpec{}, "RestartPolicy"),
    59  	}
    60  
    61  	acCmpOptions = []cmp.Option{
    62  		cmpopts.IgnoreFields(metav1.Condition{}, "LastTransitionTime"),
    63  	}
    64  )
    65  
    66  func requestWithCondition(r *autoscaling.ProvisioningRequest, conditionType string, status metav1.ConditionStatus) *autoscaling.ProvisioningRequest {
    67  	r = r.DeepCopy()
    68  	apimeta.SetStatusCondition(&r.Status.Conditions, metav1.Condition{
    69  		Type:   conditionType,
    70  		Status: status,
    71  	})
    72  	return r
    73  }
    74  
    75  func TestReconcile(t *testing.T) {
    76  	baseWorkload := utiltesting.MakeWorkload("wl", TestNamespace).
    77  		PodSets(
    78  			*utiltesting.MakePodSet("ps1", 4).
    79  				Request(corev1.ResourceCPU, "1").
    80  				Obj(),
    81  			*utiltesting.MakePodSet("ps2", 4).
    82  				Request(corev1.ResourceMemory, "1M").
    83  				Obj(),
    84  		).
    85  		ReserveQuota(utiltesting.MakeAdmission("q1").PodSets(
    86  			kueue.PodSetAssignment{
    87  				Name: "ps1",
    88  				Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
    89  					corev1.ResourceCPU: "flv1",
    90  				},
    91  				ResourceUsage: map[corev1.ResourceName]resource.Quantity{
    92  					corev1.ResourceCPU: resource.MustParse("4"),
    93  				},
    94  				Count: ptr.To[int32](4),
    95  			},
    96  			kueue.PodSetAssignment{
    97  				Name: "ps2",
    98  				Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
    99  					corev1.ResourceCPU: "flv2",
   100  				},
   101  				ResourceUsage: map[corev1.ResourceName]resource.Quantity{
   102  					corev1.ResourceCPU: resource.MustParse("3M"),
   103  				},
   104  				Count: ptr.To[int32](3),
   105  			},
   106  		).
   107  			Obj()).
   108  		AdmissionChecks(kueue.AdmissionCheckState{
   109  			Name:  "check1",
   110  			State: kueue.CheckStatePending,
   111  		}, kueue.AdmissionCheckState{
   112  			Name:  "not-provisioning",
   113  			State: kueue.CheckStatePending,
   114  		}).
   115  		Obj()
   116  
   117  	baseWorkloadWithCheck1Ready := baseWorkload.DeepCopy()
   118  	workload.SetAdmissionCheckState(&baseWorkloadWithCheck1Ready.Status.AdmissionChecks, kueue.AdmissionCheckState{
   119  		Name:  "check1",
   120  		State: kueue.CheckStateReady,
   121  	})
   122  
   123  	baseFlavor1 := utiltesting.MakeResourceFlavor("flv1").Label("f1l1", "v1").
   124  		Toleration(corev1.Toleration{
   125  			Key:      "f1t1k",
   126  			Value:    "f1t1v",
   127  			Operator: corev1.TolerationOpEqual,
   128  			Effect:   corev1.TaintEffectNoSchedule,
   129  		}).
   130  		Obj()
   131  	baseFlavor2 := utiltesting.MakeResourceFlavor("flv2").Label("f2l1", "v1").Obj()
   132  
   133  	baseRequest := &autoscaling.ProvisioningRequest{
   134  		ObjectMeta: metav1.ObjectMeta{
   135  			Namespace: TestNamespace,
   136  			Name:      "wl-check1-1",
   137  			OwnerReferences: []metav1.OwnerReference{
   138  				{
   139  					Name: "wl",
   140  				},
   141  			},
   142  		},
   143  		Spec: autoscaling.ProvisioningRequestSpec{
   144  			PodSets: []autoscaling.PodSet{
   145  				{
   146  					PodTemplateRef: autoscaling.Reference{
   147  						Name: "ppt-wl-check1-1-ps1",
   148  					},
   149  					Count: 4,
   150  				},
   151  				{
   152  					PodTemplateRef: autoscaling.Reference{
   153  						Name: "ppt-wl-check1-1-ps2",
   154  					},
   155  					Count: 3,
   156  				},
   157  			},
   158  			ProvisioningClassName: "class1",
   159  			Parameters: map[string]autoscaling.Parameter{
   160  				"p1": "v1",
   161  			},
   162  		},
   163  	}
   164  
   165  	baseTemplate1 := &corev1.PodTemplate{
   166  		ObjectMeta: metav1.ObjectMeta{
   167  			Namespace: TestNamespace,
   168  			Name:      "ppt-wl-check1-1-ps1",
   169  			OwnerReferences: []metav1.OwnerReference{
   170  				{
   171  					Name: "wl-check1-1",
   172  				},
   173  			},
   174  		},
   175  		Template: corev1.PodTemplateSpec{
   176  			Spec: corev1.PodSpec{
   177  				Containers: []corev1.Container{
   178  					{
   179  						Name: "c",
   180  						Resources: corev1.ResourceRequirements{
   181  							Requests: corev1.ResourceList{
   182  								corev1.ResourceCPU: resource.MustParse("1"),
   183  							},
   184  						},
   185  					},
   186  				},
   187  				NodeSelector: map[string]string{"f1l1": "v1"},
   188  				Tolerations: []corev1.Toleration{
   189  					{
   190  						Key:      "f1t1k",
   191  						Value:    "f1t1v",
   192  						Operator: corev1.TolerationOpEqual,
   193  						Effect:   corev1.TaintEffectNoSchedule,
   194  					},
   195  				},
   196  			},
   197  		},
   198  	}
   199  
   200  	baseTemplate2 := &corev1.PodTemplate{
   201  		ObjectMeta: metav1.ObjectMeta{
   202  			Namespace: TestNamespace,
   203  			Name:      "ppt-wl-check1-1-ps2",
   204  			OwnerReferences: []metav1.OwnerReference{
   205  				{
   206  					Name: "wl-check1-1",
   207  				},
   208  			},
   209  		},
   210  		Template: corev1.PodTemplateSpec{
   211  			Spec: corev1.PodSpec{
   212  				Containers: []corev1.Container{
   213  					{
   214  						Name: "c",
   215  						Resources: corev1.ResourceRequirements{
   216  							Requests: corev1.ResourceList{
   217  								corev1.ResourceMemory: resource.MustParse("1M"),
   218  							},
   219  						},
   220  					},
   221  				},
   222  				NodeSelector: map[string]string{"f2l1": "v1"},
   223  			},
   224  		},
   225  	}
   226  
   227  	baseConfig := &kueue.ProvisioningRequestConfig{
   228  		ObjectMeta: metav1.ObjectMeta{
   229  			Name: "config1",
   230  		},
   231  		Spec: kueue.ProvisioningRequestConfigSpec{
   232  			ProvisioningClassName: "class1",
   233  			Parameters: map[string]kueue.Parameter{
   234  				"p1": "v1",
   235  			},
   236  		},
   237  	}
   238  
   239  	baseCheck := utiltesting.MakeAdmissionCheck("check1").
   240  		ControllerName(ControllerName).
   241  		Parameters(kueue.GroupVersion.Group, ConfigKind, "config1").
   242  		Obj()
   243  
   244  	cases := map[string]struct {
   245  		requests             []autoscaling.ProvisioningRequest
   246  		templates            []corev1.PodTemplate
   247  		checks               []kueue.AdmissionCheck
   248  		configs              []kueue.ProvisioningRequestConfig
   249  		flavors              []kueue.ResourceFlavor
   250  		workload             *kueue.Workload
   251  		maxRetries           int32
   252  		wantReconcileError   error
   253  		wantWorkloads        map[string]*kueue.Workload
   254  		wantRequests         map[string]*autoscaling.ProvisioningRequest
   255  		wantTemplates        map[string]*corev1.PodTemplate
   256  		wantRequestsNotFound []string
   257  		wantEvents           []utiltesting.EventRecord
   258  	}{
   259  		"unrelated workload": {
   260  			workload: utiltesting.MakeWorkload("wl", "ns").Obj(),
   261  		},
   262  		"unrelated workload with reservation": {
   263  			workload: utiltesting.MakeWorkload("wl", "ns").
   264  				ReserveQuota(utiltesting.MakeAdmission("q1").Obj()).
   265  				Obj(),
   266  		},
   267  		"unrelated admitted workload": {
   268  			workload: utiltesting.MakeWorkload("wl", "ns").
   269  				ReserveQuota(utiltesting.MakeAdmission("q1").Obj()).
   270  				Admitted(true).
   271  				Obj(),
   272  		},
   273  		"missing config": {
   274  			workload: baseWorkload.DeepCopy(),
   275  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   276  			wantWorkloads: map[string]*kueue.Workload{
   277  				baseWorkload.Name: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   278  					AdmissionChecks(kueue.AdmissionCheckState{
   279  						Name:    "check1",
   280  						State:   kueue.CheckStatePending,
   281  						Message: CheckInactiveMessage,
   282  					}, kueue.AdmissionCheckState{
   283  						Name:  "not-provisioning",
   284  						State: kueue.CheckStatePending,
   285  					}).
   286  					Obj(),
   287  			},
   288  		},
   289  		"with config": {
   290  			workload: baseWorkload.DeepCopy(),
   291  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   292  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   293  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   294  			wantWorkloads: map[string]*kueue.Workload{
   295  				baseWorkload.Name: baseWorkload.DeepCopy(),
   296  			},
   297  			wantRequests: map[string]*autoscaling.ProvisioningRequest{
   298  				baseRequest.Name: baseRequest.DeepCopy(),
   299  			},
   300  			wantTemplates: map[string]*corev1.PodTemplate{
   301  				baseTemplate1.Name: baseTemplate1.DeepCopy(),
   302  				baseTemplate2.Name: baseTemplate2.DeepCopy(),
   303  			},
   304  			wantEvents: []utiltesting.EventRecord{
   305  				{
   306  					Key:       client.ObjectKeyFromObject(baseWorkload),
   307  					EventType: corev1.EventTypeNormal,
   308  					Reason:    "ProvisioningRequestCreated",
   309  					Message:   `Created ProvisioningRequest: "wl-check1-1"`,
   310  				},
   311  			},
   312  		},
   313  		"remove unnecessary requests": {
   314  			workload: baseWorkload.DeepCopy(),
   315  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   316  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   317  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   318  			requests: []autoscaling.ProvisioningRequest{
   319  				{
   320  					ObjectMeta: metav1.ObjectMeta{
   321  						Namespace: TestNamespace,
   322  						Name:      "wl-check2",
   323  						OwnerReferences: []metav1.OwnerReference{
   324  							{
   325  								Name: "wl",
   326  							},
   327  						},
   328  					},
   329  				},
   330  			},
   331  			wantWorkloads:        map[string]*kueue.Workload{baseWorkload.Name: baseWorkload.DeepCopy()},
   332  			wantRequestsNotFound: []string{"wl-check2"},
   333  			wantEvents: []utiltesting.EventRecord{
   334  				{
   335  					Key:       client.ObjectKeyFromObject(baseWorkload),
   336  					EventType: corev1.EventTypeNormal,
   337  					Reason:    "ProvisioningRequestCreated",
   338  					Message:   `Created ProvisioningRequest: "wl-check1-1"`,
   339  				},
   340  			},
   341  		},
   342  		"missing one template": {
   343  			workload:  baseWorkload.DeepCopy(),
   344  			checks:    []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   345  			flavors:   []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   346  			configs:   []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   347  			requests:  []autoscaling.ProvisioningRequest{*baseRequest.DeepCopy()},
   348  			templates: []corev1.PodTemplate{*baseTemplate1.DeepCopy()},
   349  			wantWorkloads: map[string]*kueue.Workload{
   350  				baseWorkload.Name: baseWorkload.DeepCopy(),
   351  			},
   352  			wantRequests: map[string]*autoscaling.ProvisioningRequest{
   353  				baseRequest.Name: baseRequest.DeepCopy(),
   354  			},
   355  			wantTemplates: map[string]*corev1.PodTemplate{
   356  				baseTemplate1.Name: baseTemplate1.DeepCopy(),
   357  				baseTemplate2.Name: baseTemplate2.DeepCopy(),
   358  			},
   359  		},
   360  		"request out of sync": {
   361  			workload: baseWorkload.DeepCopy(),
   362  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   363  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   364  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   365  			requests: []autoscaling.ProvisioningRequest{
   366  				{
   367  					ObjectMeta: metav1.ObjectMeta{
   368  						Namespace: TestNamespace,
   369  						Name:      "wl-check1-1",
   370  						OwnerReferences: []metav1.OwnerReference{
   371  							{
   372  								Name: "wl",
   373  							},
   374  						},
   375  					},
   376  					Spec: autoscaling.ProvisioningRequestSpec{
   377  						PodSets: []autoscaling.PodSet{
   378  							{
   379  								PodTemplateRef: autoscaling.Reference{
   380  									Name: "ppt-wl-check1-1-main",
   381  								},
   382  								Count: 1,
   383  							},
   384  						},
   385  						ProvisioningClassName: "class1",
   386  						Parameters: map[string]autoscaling.Parameter{
   387  							"p1": "v0",
   388  						},
   389  					},
   390  				},
   391  			},
   392  			wantWorkloads: map[string]*kueue.Workload{
   393  				baseWorkload.Name: baseWorkload.DeepCopy(),
   394  			},
   395  			wantRequests: map[string]*autoscaling.ProvisioningRequest{
   396  				baseRequest.Name: baseRequest.DeepCopy(),
   397  			},
   398  			wantTemplates: map[string]*corev1.PodTemplate{
   399  				baseTemplate1.Name: baseTemplate1.DeepCopy(),
   400  				baseTemplate2.Name: baseTemplate2.DeepCopy(),
   401  			},
   402  			wantEvents: []utiltesting.EventRecord{
   403  				{
   404  					Key:       client.ObjectKeyFromObject(baseWorkload),
   405  					EventType: corev1.EventTypeNormal,
   406  					Reason:    "ProvisioningRequestCreated",
   407  					Message:   `Created ProvisioningRequest: "wl-check1-1"`,
   408  				},
   409  			},
   410  		},
   411  		"request removed on workload finished": {
   412  			workload: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   413  				Condition(metav1.Condition{
   414  					Type:   kueue.WorkloadFinished,
   415  					Status: metav1.ConditionTrue,
   416  				}).
   417  				Obj(),
   418  
   419  			checks:               []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   420  			flavors:              []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   421  			configs:              []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   422  			requests:             []autoscaling.ProvisioningRequest{*baseRequest.DeepCopy()},
   423  			templates:            []corev1.PodTemplate{*baseTemplate1.DeepCopy(), *baseTemplate2.DeepCopy()},
   424  			wantRequestsNotFound: []string{"wl-check1"},
   425  		},
   426  		"when request fails and is retried": {
   427  			workload: baseWorkload.DeepCopy(),
   428  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   429  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   430  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   431  			requests: []autoscaling.ProvisioningRequest{
   432  				*requestWithCondition(baseRequest, autoscaling.Failed, metav1.ConditionTrue),
   433  			},
   434  			maxRetries: 2,
   435  			templates:  []corev1.PodTemplate{*baseTemplate1.DeepCopy(), *baseTemplate2.DeepCopy()},
   436  			wantWorkloads: map[string]*kueue.Workload{
   437  				baseWorkload.Name: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   438  					AdmissionChecks(kueue.AdmissionCheckState{
   439  						Name:    "check1",
   440  						State:   kueue.CheckStatePending,
   441  						Message: "Retrying after failure: ",
   442  					}, kueue.AdmissionCheckState{
   443  						Name:  "not-provisioning",
   444  						State: kueue.CheckStatePending,
   445  					}).
   446  					Obj(),
   447  			},
   448  		},
   449  		"when request fails, and there is no retry": {
   450  			workload: baseWorkload.DeepCopy(),
   451  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   452  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   453  			configs: []kueue.ProvisioningRequestConfig{
   454  				{
   455  					ObjectMeta: metav1.ObjectMeta{
   456  						Name: "config1",
   457  					},
   458  					Spec: kueue.ProvisioningRequestConfigSpec{
   459  						ProvisioningClassName: "class1",
   460  						Parameters: map[string]kueue.Parameter{
   461  							"p1": "v1",
   462  						},
   463  					},
   464  				},
   465  			},
   466  			requests: []autoscaling.ProvisioningRequest{
   467  				*requestWithCondition(baseRequest, autoscaling.Failed, metav1.ConditionTrue),
   468  			},
   469  			templates: []corev1.PodTemplate{*baseTemplate1.DeepCopy(), *baseTemplate2.DeepCopy()},
   470  			wantWorkloads: map[string]*kueue.Workload{
   471  				baseWorkload.Name: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   472  					AdmissionChecks(kueue.AdmissionCheckState{
   473  						Name:  "check1",
   474  						State: kueue.CheckStateRejected,
   475  					}, kueue.AdmissionCheckState{
   476  						Name:  "not-provisioning",
   477  						State: kueue.CheckStatePending,
   478  					}).
   479  					Obj(),
   480  			},
   481  		},
   482  		"when request is provisioned": {
   483  			workload: baseWorkload.DeepCopy(),
   484  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   485  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   486  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   487  			requests: []autoscaling.ProvisioningRequest{
   488  				*requestWithCondition(baseRequest, autoscaling.Provisioned, metav1.ConditionTrue),
   489  			},
   490  			templates: []corev1.PodTemplate{*baseTemplate1.DeepCopy(), *baseTemplate2.DeepCopy()},
   491  			wantWorkloads: map[string]*kueue.Workload{
   492  				baseWorkload.Name: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   493  					AdmissionChecks(kueue.AdmissionCheckState{
   494  						Name:  "check1",
   495  						State: kueue.CheckStateReady,
   496  						PodSetUpdates: []kueue.PodSetUpdate{
   497  							{
   498  								Name:        "ps1",
   499  								Annotations: map[string]string{"cluster-autoscaler.kubernetes.io/consume-provisioning-request": "wl-check1-1"},
   500  							},
   501  							{
   502  								Name:        "ps2",
   503  								Annotations: map[string]string{"cluster-autoscaler.kubernetes.io/consume-provisioning-request": "wl-check1-1"},
   504  							},
   505  						},
   506  					}, kueue.AdmissionCheckState{
   507  						Name:  "not-provisioning",
   508  						State: kueue.CheckStatePending,
   509  					}).
   510  					Obj(),
   511  			},
   512  		},
   513  		"when no request is needed": {
   514  			workload: baseWorkload.DeepCopy(),
   515  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   516  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   517  			configs: []kueue.ProvisioningRequestConfig{
   518  				{ObjectMeta: metav1.ObjectMeta{
   519  					Name: "config1",
   520  				},
   521  					Spec: kueue.ProvisioningRequestConfigSpec{
   522  						ProvisioningClassName: "class1",
   523  						Parameters: map[string]kueue.Parameter{
   524  							"p1": "v1",
   525  						},
   526  						ManagedResources: []corev1.ResourceName{"example.org/gpu"},
   527  					},
   528  				},
   529  			},
   530  			wantWorkloads: map[string]*kueue.Workload{
   531  				baseWorkload.Name: (&utiltesting.WorkloadWrapper{Workload: *baseWorkload.DeepCopy()}).
   532  					AdmissionChecks(kueue.AdmissionCheckState{
   533  						Name:    "check1",
   534  						State:   kueue.CheckStateReady,
   535  						Message: NoRequestNeeded,
   536  					}, kueue.AdmissionCheckState{
   537  						Name:  "not-provisioning",
   538  						State: kueue.CheckStatePending,
   539  					}).
   540  					Obj(),
   541  			},
   542  		},
   543  		"when request is needed for one PodSet": {
   544  			workload: baseWorkload.DeepCopy(),
   545  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   546  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   547  			configs: []kueue.ProvisioningRequestConfig{
   548  				{ObjectMeta: metav1.ObjectMeta{
   549  					Name: "config1",
   550  				},
   551  					Spec: kueue.ProvisioningRequestConfigSpec{
   552  						ProvisioningClassName: "class1",
   553  						Parameters: map[string]kueue.Parameter{
   554  							"p1": "v1",
   555  						},
   556  						ManagedResources: []corev1.ResourceName{corev1.ResourceMemory},
   557  					},
   558  				},
   559  			},
   560  			wantWorkloads: map[string]*kueue.Workload{
   561  				baseWorkload.Name: baseWorkload.DeepCopy(),
   562  			},
   563  			wantRequests: map[string]*autoscaling.ProvisioningRequest{
   564  				"wl-check1-1": {
   565  					Spec: autoscaling.ProvisioningRequestSpec{
   566  						PodSets: []autoscaling.PodSet{
   567  							{
   568  								PodTemplateRef: autoscaling.Reference{
   569  									Name: "ppt-wl-check1-1-ps2",
   570  								},
   571  								Count: 3,
   572  							},
   573  						},
   574  						ProvisioningClassName: "class1",
   575  						Parameters: map[string]autoscaling.Parameter{
   576  							"p1": "v1",
   577  						},
   578  					},
   579  				},
   580  			},
   581  			wantTemplates: map[string]*corev1.PodTemplate{
   582  				baseTemplate2.Name: baseTemplate2.DeepCopy(),
   583  			},
   584  			wantEvents: []utiltesting.EventRecord{
   585  				{
   586  					Key:       client.ObjectKeyFromObject(baseWorkload),
   587  					EventType: corev1.EventTypeNormal,
   588  					Reason:    "ProvisioningRequestCreated",
   589  					Message:   `Created ProvisioningRequest: "wl-check1-1"`,
   590  				},
   591  			},
   592  		},
   593  		"when the request is removed while the check is ready; don't create the ProvReq and keep Ready state": {
   594  			workload: baseWorkloadWithCheck1Ready.DeepCopy(),
   595  			checks:   []kueue.AdmissionCheck{*baseCheck.DeepCopy()},
   596  			flavors:  []kueue.ResourceFlavor{*baseFlavor1.DeepCopy(), *baseFlavor2.DeepCopy()},
   597  			configs:  []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()},
   598  			wantWorkloads: map[string]*kueue.Workload{
   599  				baseWorkload.Name: baseWorkloadWithCheck1Ready.DeepCopy(),
   600  			},
   601  			wantRequestsNotFound: []string{
   602  				GetProvisioningRequestName("wl", "check1", 1),
   603  				GetProvisioningRequestName("wl", "check2", 1),
   604  			},
   605  		},
   606  	}
   607  
   608  	for name, tc := range cases {
   609  		t.Run(name, func(t *testing.T) {
   610  			t.Cleanup(utiltesting.SetDuringTest(&MaxRetries, tc.maxRetries))
   611  
   612  			builder, ctx := getClientBuilder()
   613  
   614  			builder = builder.WithObjects(tc.workload)
   615  			builder = builder.WithStatusSubresource(tc.workload)
   616  
   617  			builder = builder.WithLists(
   618  				&autoscaling.ProvisioningRequestList{Items: tc.requests},
   619  				&corev1.PodTemplateList{Items: tc.templates},
   620  				&kueue.ProvisioningRequestConfigList{Items: tc.configs},
   621  				&kueue.AdmissionCheckList{Items: tc.checks},
   622  				&kueue.ResourceFlavorList{Items: tc.flavors},
   623  			)
   624  
   625  			k8sclient := builder.Build()
   626  			recorder := &utiltesting.EventRecorder{}
   627  			controller, err := NewController(k8sclient, recorder)
   628  			if err != nil {
   629  				t.Fatalf("Setting up the provisioning request controller: %v", err)
   630  			}
   631  
   632  			req := reconcile.Request{
   633  				NamespacedName: types.NamespacedName{
   634  					Namespace: TestNamespace,
   635  					Name:      tc.workload.Name,
   636  				},
   637  			}
   638  			_, gotReconcileError := controller.Reconcile(ctx, req)
   639  			if diff := cmp.Diff(tc.wantReconcileError, gotReconcileError); diff != "" {
   640  				t.Errorf("unexpected reconcile error (-want/+got):\n%s", diff)
   641  			}
   642  
   643  			for name, wantWl := range tc.wantWorkloads {
   644  				gotWl := &kueue.Workload{}
   645  				if err := k8sclient.Get(ctx, types.NamespacedName{Namespace: TestNamespace, Name: name}, gotWl); err != nil {
   646  					t.Errorf("unexpected error getting workload %q", name)
   647  
   648  				}
   649  
   650  				if diff := cmp.Diff(wantWl, gotWl, wlCmpOptions...); diff != "" {
   651  					t.Errorf("unexpected workload %q (-want/+got):\n%s", name, diff)
   652  				}
   653  			}
   654  
   655  			for name, wantRequest := range tc.wantRequests {
   656  				gotRequest := &autoscaling.ProvisioningRequest{}
   657  				if err := k8sclient.Get(ctx, types.NamespacedName{Namespace: TestNamespace, Name: name}, gotRequest); err != nil {
   658  					t.Errorf("unexpected error getting request %q", name)
   659  
   660  				}
   661  
   662  				if diff := cmp.Diff(wantRequest, gotRequest, reqCmpOptions...); diff != "" {
   663  					t.Errorf("unexpected request %q (-want/+got):\n%s", name, diff)
   664  				}
   665  			}
   666  
   667  			for name, wantTemplate := range tc.wantTemplates {
   668  				gotTemplate := &corev1.PodTemplate{}
   669  				if err := k8sclient.Get(ctx, types.NamespacedName{Namespace: TestNamespace, Name: name}, gotTemplate); err != nil {
   670  					t.Errorf("unexpected error getting template %q", name)
   671  
   672  				}
   673  
   674  				if diff := cmp.Diff(wantTemplate, gotTemplate, tmplCmpOptions...); diff != "" {
   675  					t.Errorf("unexpected template %q (-want/+got):\n%s", name, diff)
   676  				}
   677  			}
   678  
   679  			for _, name := range tc.wantRequestsNotFound {
   680  				gotRequest := &autoscaling.ProvisioningRequest{}
   681  				if err := k8sclient.Get(ctx, types.NamespacedName{Namespace: TestNamespace, Name: name}, gotRequest); !apierrors.IsNotFound(err) {
   682  					t.Errorf("request %q should no longer be found", name)
   683  				}
   684  			}
   685  
   686  			if diff := cmp.Diff(tc.wantEvents, recorder.RecordedEvents); diff != "" {
   687  				t.Errorf("unexpected events (-want/+got):\n%s", diff)
   688  			}
   689  		})
   690  	}
   691  
   692  }
   693  
   694  func TestActiveOrLastPRForChecks(t *testing.T) {
   695  	baseWorkload := utiltesting.MakeWorkload("wl", TestNamespace).
   696  		PodSets(
   697  			*utiltesting.MakePodSet("main", 4).
   698  				Request(corev1.ResourceCPU, "1").
   699  				Obj(),
   700  		).
   701  		ReserveQuota(utiltesting.MakeAdmission("q1").PodSets(
   702  			kueue.PodSetAssignment{
   703  				Name: "main",
   704  				Flavors: map[corev1.ResourceName]kueue.ResourceFlavorReference{
   705  					corev1.ResourceCPU: "flv1",
   706  				},
   707  				ResourceUsage: map[corev1.ResourceName]resource.Quantity{
   708  					corev1.ResourceCPU: resource.MustParse("4"),
   709  				},
   710  				Count: ptr.To[int32](4),
   711  			},
   712  		).
   713  			Obj()).
   714  		AdmissionChecks(kueue.AdmissionCheckState{
   715  			Name:  "check",
   716  			State: kueue.CheckStatePending,
   717  		}, kueue.AdmissionCheckState{
   718  			Name:  "not-provisioning",
   719  			State: kueue.CheckStatePending,
   720  		}).
   721  		Obj()
   722  	baseConfig := &kueue.ProvisioningRequestConfig{
   723  		ObjectMeta: metav1.ObjectMeta{
   724  			Name: "config1",
   725  		},
   726  		Spec: kueue.ProvisioningRequestConfigSpec{
   727  			ProvisioningClassName: "class1",
   728  			Parameters: map[string]kueue.Parameter{
   729  				"p1": "v1",
   730  			},
   731  		},
   732  	}
   733  
   734  	baseRequest := autoscaling.ProvisioningRequest{
   735  		ObjectMeta: metav1.ObjectMeta{
   736  			Namespace: TestNamespace,
   737  			Name:      "wl-check-1",
   738  			OwnerReferences: []metav1.OwnerReference{
   739  				{
   740  					Name: "wl",
   741  				},
   742  			},
   743  		},
   744  		Spec: autoscaling.ProvisioningRequestSpec{
   745  			PodSets: []autoscaling.PodSet{
   746  				{
   747  					PodTemplateRef: autoscaling.Reference{
   748  						Name: "ppt-wl-check-1-ps1",
   749  					},
   750  					Count: 4,
   751  				},
   752  			},
   753  			ProvisioningClassName: "class1",
   754  			Parameters: map[string]autoscaling.Parameter{
   755  				"p1": "v1",
   756  			},
   757  		},
   758  	}
   759  	pr1Failed := baseRequest.DeepCopy()
   760  	pr1Failed = requestWithCondition(pr1Failed, autoscaling.Failed, metav1.ConditionTrue)
   761  	pr2Created := baseRequest.DeepCopy()
   762  	pr2Created.Name = "wl-check-2"
   763  
   764  	baseCheck := utiltesting.MakeAdmissionCheck("check").
   765  		ControllerName(ControllerName).
   766  		Parameters(kueue.GroupVersion.Group, ConfigKind, "config1").
   767  		Obj()
   768  
   769  	cases := map[string]struct {
   770  		requests   []autoscaling.ProvisioningRequest
   771  		wantResult map[string]*autoscaling.ProvisioningRequest
   772  	}{
   773  		"no provisioning requests": {},
   774  		"two provisioning requests; 1 then 2": {
   775  			requests: []autoscaling.ProvisioningRequest{
   776  				*pr1Failed.DeepCopy(),
   777  				*pr2Created.DeepCopy(),
   778  			},
   779  			wantResult: map[string]*autoscaling.ProvisioningRequest{
   780  				"check": pr2Created.DeepCopy(),
   781  			},
   782  		},
   783  		"two provisioning requests; 2 then 1": {
   784  			requests: []autoscaling.ProvisioningRequest{
   785  				*pr2Created.DeepCopy(),
   786  				*pr1Failed.DeepCopy(),
   787  			},
   788  			wantResult: map[string]*autoscaling.ProvisioningRequest{
   789  				"check": pr2Created.DeepCopy(),
   790  			},
   791  		},
   792  	}
   793  
   794  	for name, tc := range cases {
   795  		t.Run(name, func(t *testing.T) {
   796  			workload := baseWorkload.DeepCopy()
   797  			relevantChecks := []string{"check"}
   798  			checks := []kueue.AdmissionCheck{*baseCheck.DeepCopy()}
   799  			configs := []kueue.ProvisioningRequestConfig{*baseConfig.DeepCopy()}
   800  
   801  			builder, ctx := getClientBuilder()
   802  
   803  			builder = builder.WithObjects(workload)
   804  			builder = builder.WithStatusSubresource(workload)
   805  
   806  			builder = builder.WithLists(
   807  				&autoscaling.ProvisioningRequestList{Items: tc.requests},
   808  				&kueue.ProvisioningRequestConfigList{Items: configs},
   809  				&kueue.AdmissionCheckList{Items: checks},
   810  			)
   811  
   812  			k8sclient := builder.Build()
   813  			recorder := &utiltesting.EventRecorder{}
   814  			controller, err := NewController(k8sclient, recorder)
   815  			if err != nil {
   816  				t.Fatalf("Setting up the provisioning request controller: %v", err)
   817  			}
   818  
   819  			gotResult := controller.activeOrLastPRForChecks(ctx, workload, relevantChecks, tc.requests)
   820  			if diff := cmp.Diff(tc.wantResult, gotResult, reqCmpOptions...); diff != "" {
   821  				t.Errorf("unexpected request %q (-want/+got):\n%s", name, diff)
   822  			}
   823  		})
   824  	}
   825  }