wrappers.go

Documentation: k8s.io/kubernetes/pkg/scheduler/testing

     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  // ResourceClaimWrapper wraps a ResourceClaim inside.
   886  type ResourceClaimWrapper struct{ resourcev1alpha2.ResourceClaim }
   887  
   888  // MakeResourceClaim creates a ResourceClaim wrapper.
   889  func MakeResourceClaim() *ResourceClaimWrapper {
   890  	return &ResourceClaimWrapper{resourcev1alpha2.ResourceClaim{}}
   891  }
   892  
   893  // FromResourceClaim creates a ResourceClaim wrapper from some existing object.
   894  func FromResourceClaim(other *resourcev1alpha2.ResourceClaim) *ResourceClaimWrapper {
   895  	return &ResourceClaimWrapper{*other.DeepCopy()}
   896  }
   897  
   898  // Obj returns the inner ResourceClaim.
   899  func (wrapper *ResourceClaimWrapper) Obj() *resourcev1alpha2.ResourceClaim {
   900  	return &wrapper.ResourceClaim
   901  }
   902  
   903  // Name sets `s` as the name of the inner object.
   904  func (wrapper *ResourceClaimWrapper) Name(s string) *ResourceClaimWrapper {
   905  	wrapper.SetName(s)
   906  	return wrapper
   907  }
   908  
   909  // UID sets `s` as the UID of the inner object.
   910  func (wrapper *ResourceClaimWrapper) UID(s string) *ResourceClaimWrapper {
   911  	wrapper.SetUID(types.UID(s))
   912  	return wrapper
   913  }
   914  
   915  // Namespace sets `s` as the namespace of the inner object.
   916  func (wrapper *ResourceClaimWrapper) Namespace(s string) *ResourceClaimWrapper {
   917  	wrapper.SetNamespace(s)
   918  	return wrapper
   919  }
   920  
   921  // OwnerReference updates the owning controller of the object.
   922  func (wrapper *ResourceClaimWrapper) OwnerReference(name, uid string, gvk schema.GroupVersionKind) *ResourceClaimWrapper {
   923  	wrapper.OwnerReferences = []metav1.OwnerReference{
   924  		{
   925  			APIVersion: gvk.GroupVersion().String(),
   926  			Kind:       gvk.Kind,
   927  			Name:       name,
   928  			UID:        types.UID(uid),
   929  			Controller: ptr.To(true),
   930  		},
   931  	}
   932  	return wrapper
   933  }
   934  
   935  // AllocationMode sets the allocation mode of the inner object.
   936  func (wrapper *ResourceClaimWrapper) AllocationMode(a resourcev1alpha2.AllocationMode) *ResourceClaimWrapper {
   937  	wrapper.ResourceClaim.Spec.AllocationMode = a
   938  	return wrapper
   939  }
   940  
   941  // ResourceClassName sets the resource class name of the inner object.
   942  func (wrapper *ResourceClaimWrapper) ResourceClassName(name string) *ResourceClaimWrapper {
   943  	wrapper.ResourceClaim.Spec.ResourceClassName = name
   944  	return wrapper
   945  }
   946  
   947  // Allocation sets the allocation of the inner object.
   948  func (wrapper *ResourceClaimWrapper) Allocation(allocation *resourcev1alpha2.AllocationResult) *ResourceClaimWrapper {
   949  	wrapper.ResourceClaim.Status.Allocation = allocation
   950  	return wrapper
   951  }
   952  
   953  // DeallocationRequested sets that field of the inner object.
   954  func (wrapper *ResourceClaimWrapper) DeallocationRequested(deallocationRequested bool) *ResourceClaimWrapper {
   955  	wrapper.ResourceClaim.Status.DeallocationRequested = deallocationRequested
   956  	return wrapper
   957  }
   958  
   959  // ReservedFor sets that field of the inner object.
   960  func (wrapper *ResourceClaimWrapper) ReservedFor(consumers ...resourcev1alpha2.ResourceClaimConsumerReference) *ResourceClaimWrapper {
   961  	wrapper.ResourceClaim.Status.ReservedFor = consumers
   962  	return wrapper
   963  }
   964  
   965  // PodSchedulingWrapper wraps a PodSchedulingContext inside.
   966  type PodSchedulingWrapper struct {
   967  	resourcev1alpha2.PodSchedulingContext
   968  }
   969  
   970  // MakePodSchedulingContexts creates a PodSchedulingContext wrapper.
   971  func MakePodSchedulingContexts() *PodSchedulingWrapper {
   972  	return &PodSchedulingWrapper{resourcev1alpha2.PodSchedulingContext{}}
   973  }
   974  
   975  // FromPodSchedulingContexts creates a PodSchedulingContext wrapper from an existing object.
   976  func FromPodSchedulingContexts(other *resourcev1alpha2.PodSchedulingContext) *PodSchedulingWrapper {
   977  	return &PodSchedulingWrapper{*other.DeepCopy()}
   978  }
   979  
   980  // Obj returns the inner object.
   981  func (wrapper *PodSchedulingWrapper) Obj() *resourcev1alpha2.PodSchedulingContext {
   982  	return &wrapper.PodSchedulingContext
   983  }
   984  
   985  // Name sets `s` as the name of the inner object.
   986  func (wrapper *PodSchedulingWrapper) Name(s string) *PodSchedulingWrapper {
   987  	wrapper.SetName(s)
   988  	return wrapper
   989  }
   990  
   991  // UID sets `s` as the UID of the inner object.
   992  func (wrapper *PodSchedulingWrapper) UID(s string) *PodSchedulingWrapper {
   993  	wrapper.SetUID(types.UID(s))
   994  	return wrapper
   995  }
   996  
   997  // Namespace sets `s` as the namespace of the inner object.
   998  func (wrapper *PodSchedulingWrapper) Namespace(s string) *PodSchedulingWrapper {
   999  	wrapper.SetNamespace(s)
  1000  	return wrapper
  1001  }
  1002  
  1003  // OwnerReference updates the owning controller of the inner object.
  1004  func (wrapper *PodSchedulingWrapper) OwnerReference(name, uid string, gvk schema.GroupVersionKind) *PodSchedulingWrapper {
  1005  	wrapper.OwnerReferences = []metav1.OwnerReference{
  1006  		{
  1007  			APIVersion:         gvk.GroupVersion().String(),
  1008  			Kind:               gvk.Kind,
  1009  			Name:               name,
  1010  			UID:                types.UID(uid),
  1011  			Controller:         ptr.To(true),
  1012  			BlockOwnerDeletion: ptr.To(true),
  1013  		},
  1014  	}
  1015  	return wrapper
  1016  }
  1017  
  1018  // Label applies a {k,v} label pair to the inner object
  1019  func (wrapper *PodSchedulingWrapper) Label(k, v string) *PodSchedulingWrapper {
  1020  	if wrapper.Labels == nil {
  1021  		wrapper.Labels = make(map[string]string)
  1022  	}
  1023  	wrapper.Labels[k] = v
  1024  	return wrapper
  1025  }
  1026  
  1027  // SelectedNode sets that field of the inner object.
  1028  func (wrapper *PodSchedulingWrapper) SelectedNode(s string) *PodSchedulingWrapper {
  1029  	wrapper.Spec.SelectedNode = s
  1030  	return wrapper
  1031  }
  1032  
  1033  // PotentialNodes sets that field of the inner object.
  1034  func (wrapper *PodSchedulingWrapper) PotentialNodes(nodes ...string) *PodSchedulingWrapper {
  1035  	wrapper.Spec.PotentialNodes = nodes
  1036  	return wrapper
  1037  }
  1038  
  1039  // ResourceClaims sets that field of the inner object.
  1040  func (wrapper *PodSchedulingWrapper) ResourceClaims(statuses ...resourcev1alpha2.ResourceClaimSchedulingStatus) *PodSchedulingWrapper {
  1041  	wrapper.Status.ResourceClaims = statuses
  1042  	return wrapper
  1043  }
  1044
View as plain text