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  }