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