k8s.io/kubernetes@v1.29.3/pkg/scheduler/testing/wrappers.go (about) 1 /* 2 Copyright 2019 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 testing 18 19 import ( 20 "fmt" 21 22 v1 "k8s.io/api/core/v1" 23 resourcev1alpha2 "k8s.io/api/resource/v1alpha2" 24 "k8s.io/apimachinery/pkg/api/resource" 25 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" 26 "k8s.io/apimachinery/pkg/runtime/schema" 27 "k8s.io/apimachinery/pkg/types" 28 imageutils "k8s.io/kubernetes/test/utils/image" 29 "k8s.io/utils/ptr" 30 ) 31 32 var zero int64 33 34 // NodeSelectorWrapper wraps a NodeSelector inside. 35 type NodeSelectorWrapper struct{ v1.NodeSelector } 36 37 // MakeNodeSelector creates a NodeSelector wrapper. 38 func MakeNodeSelector() *NodeSelectorWrapper { 39 return &NodeSelectorWrapper{v1.NodeSelector{}} 40 } 41 42 // In injects a matchExpression (with an operator IN) as a selectorTerm 43 // to the inner nodeSelector. 44 // NOTE: appended selecterTerms are ORed. 45 func (s *NodeSelectorWrapper) In(key string, vals []string) *NodeSelectorWrapper { 46 expression := v1.NodeSelectorRequirement{ 47 Key: key, 48 Operator: v1.NodeSelectorOpIn, 49 Values: vals, 50 } 51 selectorTerm := v1.NodeSelectorTerm{} 52 selectorTerm.MatchExpressions = append(selectorTerm.MatchExpressions, expression) 53 s.NodeSelectorTerms = append(s.NodeSelectorTerms, selectorTerm) 54 return s 55 } 56 57 // NotIn injects a matchExpression (with an operator NotIn) as a selectorTerm 58 // to the inner nodeSelector. 59 func (s *NodeSelectorWrapper) NotIn(key string, vals []string) *NodeSelectorWrapper { 60 expression := v1.NodeSelectorRequirement{ 61 Key: key, 62 Operator: v1.NodeSelectorOpNotIn, 63 Values: vals, 64 } 65 selectorTerm := v1.NodeSelectorTerm{} 66 selectorTerm.MatchExpressions = append(selectorTerm.MatchExpressions, expression) 67 s.NodeSelectorTerms = append(s.NodeSelectorTerms, selectorTerm) 68 return s 69 } 70 71 // Obj returns the inner NodeSelector. 72 func (s *NodeSelectorWrapper) Obj() *v1.NodeSelector { 73 return &s.NodeSelector 74 } 75 76 // LabelSelectorWrapper wraps a LabelSelector inside. 77 type LabelSelectorWrapper struct{ metav1.LabelSelector } 78 79 // MakeLabelSelector creates a LabelSelector wrapper. 80 func MakeLabelSelector() *LabelSelectorWrapper { 81 return &LabelSelectorWrapper{metav1.LabelSelector{}} 82 } 83 84 // Label applies a {k,v} pair to the inner LabelSelector. 85 func (s *LabelSelectorWrapper) Label(k, v string) *LabelSelectorWrapper { 86 if s.MatchLabels == nil { 87 s.MatchLabels = make(map[string]string) 88 } 89 s.MatchLabels[k] = v 90 return s 91 } 92 93 // In injects a matchExpression (with an operator In) to the inner labelSelector. 94 func (s *LabelSelectorWrapper) In(key string, vals []string) *LabelSelectorWrapper { 95 expression := metav1.LabelSelectorRequirement{ 96 Key: key, 97 Operator: metav1.LabelSelectorOpIn, 98 Values: vals, 99 } 100 s.MatchExpressions = append(s.MatchExpressions, expression) 101 return s 102 } 103 104 // NotIn injects a matchExpression (with an operator NotIn) to the inner labelSelector. 105 func (s *LabelSelectorWrapper) NotIn(key string, vals []string) *LabelSelectorWrapper { 106 expression := metav1.LabelSelectorRequirement{ 107 Key: key, 108 Operator: metav1.LabelSelectorOpNotIn, 109 Values: vals, 110 } 111 s.MatchExpressions = append(s.MatchExpressions, expression) 112 return s 113 } 114 115 // Exists injects a matchExpression (with an operator Exists) to the inner labelSelector. 116 func (s *LabelSelectorWrapper) Exists(k string) *LabelSelectorWrapper { 117 expression := metav1.LabelSelectorRequirement{ 118 Key: k, 119 Operator: metav1.LabelSelectorOpExists, 120 } 121 s.MatchExpressions = append(s.MatchExpressions, expression) 122 return s 123 } 124 125 // NotExist injects a matchExpression (with an operator NotExist) to the inner labelSelector. 126 func (s *LabelSelectorWrapper) NotExist(k string) *LabelSelectorWrapper { 127 expression := metav1.LabelSelectorRequirement{ 128 Key: k, 129 Operator: metav1.LabelSelectorOpDoesNotExist, 130 } 131 s.MatchExpressions = append(s.MatchExpressions, expression) 132 return s 133 } 134 135 // Obj returns the inner LabelSelector. 136 func (s *LabelSelectorWrapper) Obj() *metav1.LabelSelector { 137 return &s.LabelSelector 138 } 139 140 // ContainerWrapper wraps a Container inside. 141 type ContainerWrapper struct{ v1.Container } 142 143 // MakeContainer creates a Container wrapper. 144 func MakeContainer() *ContainerWrapper { 145 return &ContainerWrapper{v1.Container{}} 146 } 147 148 // Obj returns the inner Container. 149 func (c *ContainerWrapper) Obj() v1.Container { 150 return c.Container 151 } 152 153 // Name sets `n` as the name of the inner Container. 154 func (c *ContainerWrapper) Name(n string) *ContainerWrapper { 155 c.Container.Name = n 156 return c 157 } 158 159 // Image sets `image` as the image of the inner Container. 160 func (c *ContainerWrapper) Image(image string) *ContainerWrapper { 161 c.Container.Image = image 162 return c 163 } 164 165 // HostPort sets `hostPort` as the host port of the inner Container. 166 func (c *ContainerWrapper) HostPort(hostPort int32) *ContainerWrapper { 167 c.Container.Ports = []v1.ContainerPort{{HostPort: hostPort}} 168 return c 169 } 170 171 // ContainerPort sets `ports` as the ports of the inner Container. 172 func (c *ContainerWrapper) ContainerPort(ports []v1.ContainerPort) *ContainerWrapper { 173 c.Container.Ports = ports 174 return c 175 } 176 177 // Resources sets the container resources to the given resource map. 178 func (c *ContainerWrapper) Resources(resMap map[v1.ResourceName]string) *ContainerWrapper { 179 res := v1.ResourceList{} 180 for k, v := range resMap { 181 res[k] = resource.MustParse(v) 182 } 183 c.Container.Resources = v1.ResourceRequirements{ 184 Requests: res, 185 Limits: res, 186 } 187 return c 188 } 189 190 // ResourceRequests sets the container resources requests to the given resource map of requests. 191 func (c *ContainerWrapper) ResourceRequests(reqMap map[v1.ResourceName]string) *ContainerWrapper { 192 res := v1.ResourceList{} 193 for k, v := range reqMap { 194 res[k] = resource.MustParse(v) 195 } 196 c.Container.Resources = v1.ResourceRequirements{ 197 Requests: res, 198 } 199 return c 200 } 201 202 // ResourceLimits sets the container resource limits to the given resource map. 203 func (c *ContainerWrapper) ResourceLimits(limMap map[v1.ResourceName]string) *ContainerWrapper { 204 res := v1.ResourceList{} 205 for k, v := range limMap { 206 res[k] = resource.MustParse(v) 207 } 208 c.Container.Resources = v1.ResourceRequirements{ 209 Limits: res, 210 } 211 return c 212 } 213 214 // PodWrapper wraps a Pod inside. 215 type PodWrapper struct{ v1.Pod } 216 217 // MakePod creates a Pod wrapper. 218 func MakePod() *PodWrapper { 219 return &PodWrapper{v1.Pod{}} 220 } 221 222 // Obj returns the inner Pod. 223 func (p *PodWrapper) Obj() *v1.Pod { 224 return &p.Pod 225 } 226 227 // Name sets `s` as the name of the inner pod. 228 func (p *PodWrapper) Name(s string) *PodWrapper { 229 p.SetName(s) 230 return p 231 } 232 233 // UID sets `s` as the UID of the inner pod. 234 func (p *PodWrapper) UID(s string) *PodWrapper { 235 p.SetUID(types.UID(s)) 236 return p 237 } 238 239 // SchedulerName sets `s` as the scheduler name of the inner pod. 240 func (p *PodWrapper) SchedulerName(s string) *PodWrapper { 241 p.Spec.SchedulerName = s 242 return p 243 } 244 245 // Namespace sets `s` as the namespace of the inner pod. 246 func (p *PodWrapper) Namespace(s string) *PodWrapper { 247 p.SetNamespace(s) 248 return p 249 } 250 251 // OwnerReference updates the owning controller of the pod. 252 func (p *PodWrapper) OwnerReference(name string, gvk schema.GroupVersionKind) *PodWrapper { 253 p.OwnerReferences = []metav1.OwnerReference{ 254 { 255 APIVersion: gvk.GroupVersion().String(), 256 Kind: gvk.Kind, 257 Name: name, 258 Controller: ptr.To(true), 259 }, 260 } 261 return p 262 } 263 264 // Container appends a container into PodSpec of the inner pod. 265 func (p *PodWrapper) Container(s string) *PodWrapper { 266 name := fmt.Sprintf("con%d", len(p.Spec.Containers)) 267 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name(name).Image(s).Obj()) 268 return p 269 } 270 271 // Containers sets `containers` to the PodSpec of the inner pod. 272 func (p *PodWrapper) Containers(containers []v1.Container) *PodWrapper { 273 p.Spec.Containers = containers 274 return p 275 } 276 277 // PodResourceClaims appends PodResourceClaims into PodSpec of the inner pod. 278 func (p *PodWrapper) PodResourceClaims(podResourceClaims ...v1.PodResourceClaim) *PodWrapper { 279 p.Spec.ResourceClaims = append(p.Spec.ResourceClaims, podResourceClaims...) 280 return p 281 } 282 283 // Priority sets a priority value into PodSpec of the inner pod. 284 func (p *PodWrapper) Priority(val int32) *PodWrapper { 285 p.Spec.Priority = &val 286 return p 287 } 288 289 // CreationTimestamp sets the inner pod's CreationTimestamp. 290 func (p *PodWrapper) CreationTimestamp(t metav1.Time) *PodWrapper { 291 p.ObjectMeta.CreationTimestamp = t 292 return p 293 } 294 295 // Terminating sets the inner pod's deletionTimestamp to current timestamp. 296 func (p *PodWrapper) Terminating() *PodWrapper { 297 now := metav1.Now() 298 p.DeletionTimestamp = &now 299 return p 300 } 301 302 // ZeroTerminationGracePeriod sets the TerminationGracePeriodSeconds of the inner pod to zero. 303 func (p *PodWrapper) ZeroTerminationGracePeriod() *PodWrapper { 304 p.Spec.TerminationGracePeriodSeconds = &zero 305 return p 306 } 307 308 // Node sets `s` as the nodeName of the inner pod. 309 func (p *PodWrapper) Node(s string) *PodWrapper { 310 p.Spec.NodeName = s 311 return p 312 } 313 314 // NodeSelector sets `m` as the nodeSelector of the inner pod. 315 func (p *PodWrapper) NodeSelector(m map[string]string) *PodWrapper { 316 p.Spec.NodeSelector = m 317 return p 318 } 319 320 // NodeAffinityIn creates a HARD node affinity (with the operator In) 321 // and injects into the inner pod. 322 func (p *PodWrapper) NodeAffinityIn(key string, vals []string) *PodWrapper { 323 if p.Spec.Affinity == nil { 324 p.Spec.Affinity = &v1.Affinity{} 325 } 326 if p.Spec.Affinity.NodeAffinity == nil { 327 p.Spec.Affinity.NodeAffinity = &v1.NodeAffinity{} 328 } 329 nodeSelector := MakeNodeSelector().In(key, vals).Obj() 330 p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector 331 return p 332 } 333 334 // NodeAffinityNotIn creates a HARD node affinity (with the operator NotIn) 335 // and injects into the inner pod. 336 func (p *PodWrapper) NodeAffinityNotIn(key string, vals []string) *PodWrapper { 337 if p.Spec.Affinity == nil { 338 p.Spec.Affinity = &v1.Affinity{} 339 } 340 if p.Spec.Affinity.NodeAffinity == nil { 341 p.Spec.Affinity.NodeAffinity = &v1.NodeAffinity{} 342 } 343 nodeSelector := MakeNodeSelector().NotIn(key, vals).Obj() 344 p.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = nodeSelector 345 return p 346 } 347 348 // StartTime sets `t` as .status.startTime for the inner pod. 349 func (p *PodWrapper) StartTime(t metav1.Time) *PodWrapper { 350 p.Status.StartTime = &t 351 return p 352 } 353 354 // NominatedNodeName sets `n` as the .Status.NominatedNodeName of the inner pod. 355 func (p *PodWrapper) NominatedNodeName(n string) *PodWrapper { 356 p.Status.NominatedNodeName = n 357 return p 358 } 359 360 // Phase sets `phase` as .status.Phase of the inner pod. 361 func (p *PodWrapper) Phase(phase v1.PodPhase) *PodWrapper { 362 p.Status.Phase = phase 363 return p 364 } 365 366 // Condition adds a `condition(Type, Status, Reason)` to .Status.Conditions. 367 func (p *PodWrapper) Condition(t v1.PodConditionType, s v1.ConditionStatus, r string) *PodWrapper { 368 p.Status.Conditions = append(p.Status.Conditions, v1.PodCondition{Type: t, Status: s, Reason: r}) 369 return p 370 } 371 372 // Conditions sets `conditions` as .status.Conditions of the inner pod. 373 func (p *PodWrapper) Conditions(conditions []v1.PodCondition) *PodWrapper { 374 p.Status.Conditions = append(p.Status.Conditions, conditions...) 375 return p 376 } 377 378 // Toleration creates a toleration (with the operator Exists) 379 // and injects into the inner pod. 380 func (p *PodWrapper) Toleration(key string) *PodWrapper { 381 p.Spec.Tolerations = append(p.Spec.Tolerations, v1.Toleration{ 382 Key: key, 383 Operator: v1.TolerationOpExists, 384 }) 385 return p 386 } 387 388 // HostPort creates a container with a hostPort valued `hostPort`, 389 // and injects into the inner pod. 390 func (p *PodWrapper) HostPort(port int32) *PodWrapper { 391 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name("container").Image("pause").HostPort(port).Obj()) 392 return p 393 } 394 395 // ContainerPort creates a container with ports valued `ports`, 396 // and injects into the inner pod. 397 func (p *PodWrapper) ContainerPort(ports []v1.ContainerPort) *PodWrapper { 398 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name("container").Image("pause").ContainerPort(ports).Obj()) 399 return p 400 } 401 402 // PVC creates a Volume with a PVC and injects into the inner pod. 403 func (p *PodWrapper) PVC(name string) *PodWrapper { 404 p.Spec.Volumes = append(p.Spec.Volumes, v1.Volume{ 405 Name: name, 406 VolumeSource: v1.VolumeSource{ 407 PersistentVolumeClaim: &v1.PersistentVolumeClaimVolumeSource{ClaimName: name}, 408 }, 409 }) 410 return p 411 } 412 413 // Volume creates volume and injects into the inner pod. 414 func (p *PodWrapper) Volume(volume v1.Volume) *PodWrapper { 415 p.Spec.Volumes = append(p.Spec.Volumes, volume) 416 return p 417 } 418 419 // Volumes set the volumes and inject into the inner pod. 420 func (p *PodWrapper) Volumes(volumes []v1.Volume) *PodWrapper { 421 p.Spec.Volumes = volumes 422 return p 423 } 424 425 // SchedulingGates sets `gates` as additional SchedulerGates of the inner pod. 426 func (p *PodWrapper) SchedulingGates(gates []string) *PodWrapper { 427 for _, gate := range gates { 428 p.Spec.SchedulingGates = append(p.Spec.SchedulingGates, v1.PodSchedulingGate{Name: gate}) 429 } 430 return p 431 } 432 433 // PodAffinityKind represents different kinds of PodAffinity. 434 type PodAffinityKind int 435 436 const ( 437 // NilPodAffinity is a no-op which doesn't apply any PodAffinity. 438 NilPodAffinity PodAffinityKind = iota 439 // PodAffinityWithRequiredReq applies a HARD requirement to pod.spec.affinity.PodAffinity. 440 PodAffinityWithRequiredReq 441 // PodAffinityWithPreferredReq applies a SOFT requirement to pod.spec.affinity.PodAffinity. 442 PodAffinityWithPreferredReq 443 // PodAffinityWithRequiredPreferredReq applies HARD and SOFT requirements to pod.spec.affinity.PodAffinity. 444 PodAffinityWithRequiredPreferredReq 445 // PodAntiAffinityWithRequiredReq applies a HARD requirement to pod.spec.affinity.PodAntiAffinity. 446 PodAntiAffinityWithRequiredReq 447 // PodAntiAffinityWithPreferredReq applies a SOFT requirement to pod.spec.affinity.PodAntiAffinity. 448 PodAntiAffinityWithPreferredReq 449 // PodAntiAffinityWithRequiredPreferredReq applies HARD and SOFT requirements to pod.spec.affinity.PodAntiAffinity. 450 PodAntiAffinityWithRequiredPreferredReq 451 ) 452 453 // PodAffinity creates a PodAffinity with topology key and label selector 454 // and injects into the inner pod. 455 func (p *PodWrapper) PodAffinity(topologyKey string, labelSelector *metav1.LabelSelector, kind PodAffinityKind) *PodWrapper { 456 if kind == NilPodAffinity { 457 return p 458 } 459 460 if p.Spec.Affinity == nil { 461 p.Spec.Affinity = &v1.Affinity{} 462 } 463 if p.Spec.Affinity.PodAffinity == nil { 464 p.Spec.Affinity.PodAffinity = &v1.PodAffinity{} 465 } 466 term := v1.PodAffinityTerm{LabelSelector: labelSelector, TopologyKey: topologyKey} 467 switch kind { 468 case PodAffinityWithRequiredReq: 469 p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append( 470 p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution, 471 term, 472 ) 473 case PodAffinityWithPreferredReq: 474 p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append( 475 p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 476 v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term}, 477 ) 478 case PodAffinityWithRequiredPreferredReq: 479 p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append( 480 p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution, 481 term, 482 ) 483 p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append( 484 p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 485 v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term}, 486 ) 487 } 488 return p 489 } 490 491 // PodAntiAffinity creates a PodAntiAffinity with topology key and label selector 492 // and injects into the inner pod. 493 func (p *PodWrapper) PodAntiAffinity(topologyKey string, labelSelector *metav1.LabelSelector, kind PodAffinityKind) *PodWrapper { 494 if kind == NilPodAffinity { 495 return p 496 } 497 498 if p.Spec.Affinity == nil { 499 p.Spec.Affinity = &v1.Affinity{} 500 } 501 if p.Spec.Affinity.PodAntiAffinity == nil { 502 p.Spec.Affinity.PodAntiAffinity = &v1.PodAntiAffinity{} 503 } 504 term := v1.PodAffinityTerm{LabelSelector: labelSelector, TopologyKey: topologyKey} 505 switch kind { 506 case PodAntiAffinityWithRequiredReq: 507 p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append( 508 p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution, 509 term, 510 ) 511 case PodAntiAffinityWithPreferredReq: 512 p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append( 513 p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 514 v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term}, 515 ) 516 case PodAntiAffinityWithRequiredPreferredReq: 517 p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append( 518 p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution, 519 term, 520 ) 521 p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append( 522 p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution, 523 v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term}, 524 ) 525 } 526 return p 527 } 528 529 // PodAffinityExists creates a PodAffinity with the operator "Exists" 530 // and injects into the inner pod. 531 func (p *PodWrapper) PodAffinityExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper { 532 labelSelector := MakeLabelSelector().Exists(labelKey).Obj() 533 p.PodAffinity(topologyKey, labelSelector, kind) 534 return p 535 } 536 537 // PodAntiAffinityExists creates a PodAntiAffinity with the operator "Exists" 538 // and injects into the inner pod. 539 func (p *PodWrapper) PodAntiAffinityExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper { 540 labelSelector := MakeLabelSelector().Exists(labelKey).Obj() 541 p.PodAntiAffinity(topologyKey, labelSelector, kind) 542 return p 543 } 544 545 // PodAffinityNotExists creates a PodAffinity with the operator "NotExists" 546 // and injects into the inner pod. 547 func (p *PodWrapper) PodAffinityNotExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper { 548 labelSelector := MakeLabelSelector().NotExist(labelKey).Obj() 549 p.PodAffinity(topologyKey, labelSelector, kind) 550 return p 551 } 552 553 // PodAntiAffinityNotExists creates a PodAntiAffinity with the operator "NotExists" 554 // and injects into the inner pod. 555 func (p *PodWrapper) PodAntiAffinityNotExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper { 556 labelSelector := MakeLabelSelector().NotExist(labelKey).Obj() 557 p.PodAntiAffinity(topologyKey, labelSelector, kind) 558 return p 559 } 560 561 // PodAffinityIn creates a PodAffinity with the operator "In" 562 // and injects into the inner pod. 563 func (p *PodWrapper) PodAffinityIn(labelKey, topologyKey string, vals []string, kind PodAffinityKind) *PodWrapper { 564 labelSelector := MakeLabelSelector().In(labelKey, vals).Obj() 565 p.PodAffinity(topologyKey, labelSelector, kind) 566 return p 567 } 568 569 // PodAntiAffinityIn creates a PodAntiAffinity with the operator "In" 570 // and injects into the inner pod. 571 func (p *PodWrapper) PodAntiAffinityIn(labelKey, topologyKey string, vals []string, kind PodAffinityKind) *PodWrapper { 572 labelSelector := MakeLabelSelector().In(labelKey, vals).Obj() 573 p.PodAntiAffinity(topologyKey, labelSelector, kind) 574 return p 575 } 576 577 // PodAffinityNotIn creates a PodAffinity with the operator "NotIn" 578 // and injects into the inner pod. 579 func (p *PodWrapper) PodAffinityNotIn(labelKey, topologyKey string, vals []string, kind PodAffinityKind) *PodWrapper { 580 labelSelector := MakeLabelSelector().NotIn(labelKey, vals).Obj() 581 p.PodAffinity(topologyKey, labelSelector, kind) 582 return p 583 } 584 585 // PodAntiAffinityNotIn creates a PodAntiAffinity with the operator "NotIn" 586 // and injects into the inner pod. 587 func (p *PodWrapper) PodAntiAffinityNotIn(labelKey, topologyKey string, vals []string, kind PodAffinityKind) *PodWrapper { 588 labelSelector := MakeLabelSelector().NotIn(labelKey, vals).Obj() 589 p.PodAntiAffinity(topologyKey, labelSelector, kind) 590 return p 591 } 592 593 // SpreadConstraint constructs a TopologySpreadConstraint object and injects 594 // into the inner pod. 595 func (p *PodWrapper) SpreadConstraint(maxSkew int, tpKey string, mode v1.UnsatisfiableConstraintAction, selector *metav1.LabelSelector, minDomains *int32, nodeAffinityPolicy, nodeTaintsPolicy *v1.NodeInclusionPolicy, matchLabelKeys []string) *PodWrapper { 596 c := v1.TopologySpreadConstraint{ 597 MaxSkew: int32(maxSkew), 598 TopologyKey: tpKey, 599 WhenUnsatisfiable: mode, 600 LabelSelector: selector, 601 MinDomains: minDomains, 602 NodeAffinityPolicy: nodeAffinityPolicy, 603 NodeTaintsPolicy: nodeTaintsPolicy, 604 MatchLabelKeys: matchLabelKeys, 605 } 606 p.Spec.TopologySpreadConstraints = append(p.Spec.TopologySpreadConstraints, c) 607 return p 608 } 609 610 // Label sets a {k,v} pair to the inner pod label. 611 func (p *PodWrapper) Label(k, v string) *PodWrapper { 612 if p.ObjectMeta.Labels == nil { 613 p.ObjectMeta.Labels = make(map[string]string) 614 } 615 p.ObjectMeta.Labels[k] = v 616 return p 617 } 618 619 // Labels sets all {k,v} pair provided by `labels` to the inner pod labels. 620 func (p *PodWrapper) Labels(labels map[string]string) *PodWrapper { 621 for k, v := range labels { 622 p.Label(k, v) 623 } 624 return p 625 } 626 627 // Annotation sets a {k,v} pair to the inner pod annotation. 628 func (p *PodWrapper) Annotation(key, value string) *PodWrapper { 629 metav1.SetMetaDataAnnotation(&p.ObjectMeta, key, value) 630 return p 631 } 632 633 // Annotations sets all {k,v} pair provided by `annotations` to the inner pod annotations. 634 func (p *PodWrapper) Annotations(annotations map[string]string) *PodWrapper { 635 for k, v := range annotations { 636 p.Annotation(k, v) 637 } 638 return p 639 } 640 641 // Res adds a new container to the inner pod with given resource map. 642 func (p *PodWrapper) Res(resMap map[v1.ResourceName]string) *PodWrapper { 643 if len(resMap) == 0 { 644 return p 645 } 646 647 name := fmt.Sprintf("con%d", len(p.Spec.Containers)) 648 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name(name).Image(imageutils.GetPauseImageName()).Resources(resMap).Obj()) 649 return p 650 } 651 652 // Req adds a new container to the inner pod with given resource map of requests. 653 func (p *PodWrapper) Req(reqMap map[v1.ResourceName]string) *PodWrapper { 654 if len(reqMap) == 0 { 655 return p 656 } 657 658 name := fmt.Sprintf("con%d", len(p.Spec.Containers)) 659 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name(name).Image(imageutils.GetPauseImageName()).ResourceRequests(reqMap).Obj()) 660 return p 661 } 662 663 // Lim adds a new container to the inner pod with given resource map of limits. 664 func (p *PodWrapper) Lim(limMap map[v1.ResourceName]string) *PodWrapper { 665 if len(limMap) == 0 { 666 return p 667 } 668 669 name := fmt.Sprintf("con%d", len(p.Spec.Containers)) 670 p.Spec.Containers = append(p.Spec.Containers, MakeContainer().Name(name).Image(imageutils.GetPauseImageName()).ResourceLimits(limMap).Obj()) 671 return p 672 } 673 674 // InitReq adds a new init container to the inner pod with given resource map. 675 func (p *PodWrapper) InitReq(resMap map[v1.ResourceName]string) *PodWrapper { 676 if len(resMap) == 0 { 677 return p 678 } 679 680 name := fmt.Sprintf("init-con%d", len(p.Spec.InitContainers)) 681 p.Spec.InitContainers = append(p.Spec.InitContainers, MakeContainer().Name(name).Image(imageutils.GetPauseImageName()).Resources(resMap).Obj()) 682 return p 683 } 684 685 // PreemptionPolicy sets the give preemption policy to the inner pod. 686 func (p *PodWrapper) PreemptionPolicy(policy v1.PreemptionPolicy) *PodWrapper { 687 p.Spec.PreemptionPolicy = &policy 688 return p 689 } 690 691 // Overhead sets the give ResourceList to the inner pod 692 func (p *PodWrapper) Overhead(rl v1.ResourceList) *PodWrapper { 693 p.Spec.Overhead = rl 694 return p 695 } 696 697 // NodeWrapper wraps a Node inside. 698 type NodeWrapper struct{ v1.Node } 699 700 // MakeNode creates a Node wrapper. 701 func MakeNode() *NodeWrapper { 702 w := &NodeWrapper{v1.Node{}} 703 return w.Capacity(nil) 704 } 705 706 // Obj returns the inner Node. 707 func (n *NodeWrapper) Obj() *v1.Node { 708 return &n.Node 709 } 710 711 // Name sets `s` as the name of the inner pod. 712 func (n *NodeWrapper) Name(s string) *NodeWrapper { 713 n.SetName(s) 714 return n 715 } 716 717 // UID sets `s` as the UID of the inner pod. 718 func (n *NodeWrapper) UID(s string) *NodeWrapper { 719 n.SetUID(types.UID(s)) 720 return n 721 } 722 723 // Label applies a {k,v} label pair to the inner node. 724 func (n *NodeWrapper) Label(k, v string) *NodeWrapper { 725 if n.Labels == nil { 726 n.Labels = make(map[string]string) 727 } 728 n.Labels[k] = v 729 return n 730 } 731 732 // Capacity sets the capacity and the allocatable resources of the inner node. 733 // Each entry in `resources` corresponds to a resource name and its quantity. 734 // By default, the capacity and allocatable number of pods are set to 32. 735 func (n *NodeWrapper) Capacity(resources map[v1.ResourceName]string) *NodeWrapper { 736 res := v1.ResourceList{ 737 v1.ResourcePods: resource.MustParse("32"), 738 } 739 for name, value := range resources { 740 res[name] = resource.MustParse(value) 741 } 742 n.Status.Capacity, n.Status.Allocatable = res, res 743 return n 744 } 745 746 // Images sets the images of the inner node. Each entry in `images` corresponds 747 // to an image name and its size in bytes. 748 func (n *NodeWrapper) Images(images map[string]int64) *NodeWrapper { 749 var containerImages []v1.ContainerImage 750 for name, size := range images { 751 containerImages = append(containerImages, v1.ContainerImage{Names: []string{name}, SizeBytes: size}) 752 } 753 n.Status.Images = containerImages 754 return n 755 } 756 757 // Taints applies taints to the inner node. 758 func (n *NodeWrapper) Taints(taints []v1.Taint) *NodeWrapper { 759 n.Spec.Taints = taints 760 return n 761 } 762 763 // PersistentVolumeClaimWrapper wraps a PersistentVolumeClaim inside. 764 type PersistentVolumeClaimWrapper struct{ v1.PersistentVolumeClaim } 765 766 // MakePersistentVolumeClaim creates a PersistentVolumeClaim wrapper. 767 func MakePersistentVolumeClaim() *PersistentVolumeClaimWrapper { 768 return &PersistentVolumeClaimWrapper{} 769 } 770 771 // Obj returns the inner PersistentVolumeClaim. 772 func (p *PersistentVolumeClaimWrapper) Obj() *v1.PersistentVolumeClaim { 773 return &p.PersistentVolumeClaim 774 } 775 776 // Name sets `s` as the name of the inner PersistentVolumeClaim. 777 func (p *PersistentVolumeClaimWrapper) Name(s string) *PersistentVolumeClaimWrapper { 778 p.SetName(s) 779 return p 780 } 781 782 // Namespace sets `s` as the namespace of the inner PersistentVolumeClaim. 783 func (p *PersistentVolumeClaimWrapper) Namespace(s string) *PersistentVolumeClaimWrapper { 784 p.SetNamespace(s) 785 return p 786 } 787 788 // Annotation sets a {k,v} pair to the inner PersistentVolumeClaim. 789 func (p *PersistentVolumeClaimWrapper) Annotation(key, value string) *PersistentVolumeClaimWrapper { 790 metav1.SetMetaDataAnnotation(&p.ObjectMeta, key, value) 791 return p 792 } 793 794 // VolumeName sets `name` as the volume name of the inner 795 // PersistentVolumeClaim. 796 func (p *PersistentVolumeClaimWrapper) VolumeName(name string) *PersistentVolumeClaimWrapper { 797 p.PersistentVolumeClaim.Spec.VolumeName = name 798 return p 799 } 800 801 // AccessModes sets `accessModes` as the access modes of the inner 802 // PersistentVolumeClaim. 803 func (p *PersistentVolumeClaimWrapper) AccessModes(accessModes []v1.PersistentVolumeAccessMode) *PersistentVolumeClaimWrapper { 804 p.PersistentVolumeClaim.Spec.AccessModes = accessModes 805 return p 806 } 807 808 // Resources sets `resources` as the resource requirements of the inner 809 // PersistentVolumeClaim. 810 func (p *PersistentVolumeClaimWrapper) Resources(resources v1.VolumeResourceRequirements) *PersistentVolumeClaimWrapper { 811 p.PersistentVolumeClaim.Spec.Resources = resources 812 return p 813 } 814 815 // PersistentVolumeWrapper wraps a PersistentVolume inside. 816 type PersistentVolumeWrapper struct{ v1.PersistentVolume } 817 818 // MakePersistentVolume creates a PersistentVolume wrapper. 819 func MakePersistentVolume() *PersistentVolumeWrapper { 820 return &PersistentVolumeWrapper{} 821 } 822 823 // Obj returns the inner PersistentVolume. 824 func (p *PersistentVolumeWrapper) Obj() *v1.PersistentVolume { 825 return &p.PersistentVolume 826 } 827 828 // Name sets `s` as the name of the inner PersistentVolume. 829 func (p *PersistentVolumeWrapper) Name(s string) *PersistentVolumeWrapper { 830 p.SetName(s) 831 return p 832 } 833 834 // AccessModes sets `accessModes` as the access modes of the inner 835 // PersistentVolume. 836 func (p *PersistentVolumeWrapper) AccessModes(accessModes []v1.PersistentVolumeAccessMode) *PersistentVolumeWrapper { 837 p.PersistentVolume.Spec.AccessModes = accessModes 838 return p 839 } 840 841 // Capacity sets `capacity` as the resource list of the inner PersistentVolume. 842 func (p *PersistentVolumeWrapper) Capacity(capacity v1.ResourceList) *PersistentVolumeWrapper { 843 p.PersistentVolume.Spec.Capacity = capacity 844 return p 845 } 846 847 // HostPathVolumeSource sets `src` as the host path volume source of the inner 848 // PersistentVolume. 849 func (p *PersistentVolumeWrapper) HostPathVolumeSource(src *v1.HostPathVolumeSource) *PersistentVolumeWrapper { 850 p.PersistentVolume.Spec.HostPath = src 851 return p 852 } 853 854 // ResourceClaimWrapper wraps a ResourceClaim inside. 855 type ResourceClaimWrapper struct{ resourcev1alpha2.ResourceClaim } 856 857 // MakeResourceClaim creates a ResourceClaim wrapper. 858 func MakeResourceClaim() *ResourceClaimWrapper { 859 return &ResourceClaimWrapper{resourcev1alpha2.ResourceClaim{}} 860 } 861 862 // FromResourceClaim creates a ResourceClaim wrapper from some existing object. 863 func FromResourceClaim(other *resourcev1alpha2.ResourceClaim) *ResourceClaimWrapper { 864 return &ResourceClaimWrapper{*other.DeepCopy()} 865 } 866 867 // Obj returns the inner ResourceClaim. 868 func (wrapper *ResourceClaimWrapper) Obj() *resourcev1alpha2.ResourceClaim { 869 return &wrapper.ResourceClaim 870 } 871 872 // Name sets `s` as the name of the inner object. 873 func (wrapper *ResourceClaimWrapper) Name(s string) *ResourceClaimWrapper { 874 wrapper.SetName(s) 875 return wrapper 876 } 877 878 // UID sets `s` as the UID of the inner object. 879 func (wrapper *ResourceClaimWrapper) UID(s string) *ResourceClaimWrapper { 880 wrapper.SetUID(types.UID(s)) 881 return wrapper 882 } 883 884 // Namespace sets `s` as the namespace of the inner object. 885 func (wrapper *ResourceClaimWrapper) Namespace(s string) *ResourceClaimWrapper { 886 wrapper.SetNamespace(s) 887 return wrapper 888 } 889 890 // OwnerReference updates the owning controller of the object. 891 func (wrapper *ResourceClaimWrapper) OwnerReference(name, uid string, gvk schema.GroupVersionKind) *ResourceClaimWrapper { 892 wrapper.OwnerReferences = []metav1.OwnerReference{ 893 { 894 APIVersion: gvk.GroupVersion().String(), 895 Kind: gvk.Kind, 896 Name: name, 897 UID: types.UID(uid), 898 Controller: ptr.To(true), 899 }, 900 } 901 return wrapper 902 } 903 904 // AllocationMode sets the allocation mode of the inner object. 905 func (wrapper *ResourceClaimWrapper) AllocationMode(a resourcev1alpha2.AllocationMode) *ResourceClaimWrapper { 906 wrapper.ResourceClaim.Spec.AllocationMode = a 907 return wrapper 908 } 909 910 // ResourceClassName sets the resource class name of the inner object. 911 func (wrapper *ResourceClaimWrapper) ResourceClassName(name string) *ResourceClaimWrapper { 912 wrapper.ResourceClaim.Spec.ResourceClassName = name 913 return wrapper 914 } 915 916 // Allocation sets the allocation of the inner object. 917 func (wrapper *ResourceClaimWrapper) Allocation(allocation *resourcev1alpha2.AllocationResult) *ResourceClaimWrapper { 918 wrapper.ResourceClaim.Status.Allocation = allocation 919 return wrapper 920 } 921 922 // DeallocationRequested sets that field of the inner object. 923 func (wrapper *ResourceClaimWrapper) DeallocationRequested(deallocationRequested bool) *ResourceClaimWrapper { 924 wrapper.ResourceClaim.Status.DeallocationRequested = deallocationRequested 925 return wrapper 926 } 927 928 // ReservedFor sets that field of the inner object. 929 func (wrapper *ResourceClaimWrapper) ReservedFor(consumers ...resourcev1alpha2.ResourceClaimConsumerReference) *ResourceClaimWrapper { 930 wrapper.ResourceClaim.Status.ReservedFor = consumers 931 return wrapper 932 } 933 934 // PodSchedulingWrapper wraps a PodSchedulingContext inside. 935 type PodSchedulingWrapper struct { 936 resourcev1alpha2.PodSchedulingContext 937 } 938 939 // MakePodSchedulingContexts creates a PodSchedulingContext wrapper. 940 func MakePodSchedulingContexts() *PodSchedulingWrapper { 941 return &PodSchedulingWrapper{resourcev1alpha2.PodSchedulingContext{}} 942 } 943 944 // FromPodSchedulingContexts creates a PodSchedulingContext wrapper from an existing object. 945 func FromPodSchedulingContexts(other *resourcev1alpha2.PodSchedulingContext) *PodSchedulingWrapper { 946 return &PodSchedulingWrapper{*other.DeepCopy()} 947 } 948 949 // Obj returns the inner object. 950 func (wrapper *PodSchedulingWrapper) Obj() *resourcev1alpha2.PodSchedulingContext { 951 return &wrapper.PodSchedulingContext 952 } 953 954 // Name sets `s` as the name of the inner object. 955 func (wrapper *PodSchedulingWrapper) Name(s string) *PodSchedulingWrapper { 956 wrapper.SetName(s) 957 return wrapper 958 } 959 960 // UID sets `s` as the UID of the inner object. 961 func (wrapper *PodSchedulingWrapper) UID(s string) *PodSchedulingWrapper { 962 wrapper.SetUID(types.UID(s)) 963 return wrapper 964 } 965 966 // Namespace sets `s` as the namespace of the inner object. 967 func (wrapper *PodSchedulingWrapper) Namespace(s string) *PodSchedulingWrapper { 968 wrapper.SetNamespace(s) 969 return wrapper 970 } 971 972 // OwnerReference updates the owning controller of the inner object. 973 func (wrapper *PodSchedulingWrapper) OwnerReference(name, uid string, gvk schema.GroupVersionKind) *PodSchedulingWrapper { 974 wrapper.OwnerReferences = []metav1.OwnerReference{ 975 { 976 APIVersion: gvk.GroupVersion().String(), 977 Kind: gvk.Kind, 978 Name: name, 979 UID: types.UID(uid), 980 Controller: ptr.To(true), 981 BlockOwnerDeletion: ptr.To(true), 982 }, 983 } 984 return wrapper 985 } 986 987 // Label applies a {k,v} label pair to the inner object 988 func (wrapper *PodSchedulingWrapper) Label(k, v string) *PodSchedulingWrapper { 989 if wrapper.Labels == nil { 990 wrapper.Labels = make(map[string]string) 991 } 992 wrapper.Labels[k] = v 993 return wrapper 994 } 995 996 // SelectedNode sets that field of the inner object. 997 func (wrapper *PodSchedulingWrapper) SelectedNode(s string) *PodSchedulingWrapper { 998 wrapper.Spec.SelectedNode = s 999 return wrapper 1000 } 1001 1002 // PotentialNodes sets that field of the inner object. 1003 func (wrapper *PodSchedulingWrapper) PotentialNodes(nodes ...string) *PodSchedulingWrapper { 1004 wrapper.Spec.PotentialNodes = nodes 1005 return wrapper 1006 } 1007 1008 // ResourceClaims sets that field of the inner object. 1009 func (wrapper *PodSchedulingWrapper) ResourceClaims(statuses ...resourcev1alpha2.ResourceClaimSchedulingStatus) *PodSchedulingWrapper { 1010 wrapper.Status.ResourceClaims = statuses 1011 return wrapper 1012 }