generate.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/testutils

     1  // Copyright 2020 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package testutils
    15  
    16  import (
    17  	"fmt"
    18  
    19  	v1 "k8s.io/api/core/v1"
    20  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    21  	"k8s.io/apimachinery/pkg/runtime"
    22  )
    23  
    24  // PodArg by default use `Status=corev1.PodRunning` and `Namespace=metav1.NamespaceDefault`.
    25  // For the others, the default values are empty.
    26  type PodArg struct {
    27  	Name            string
    28  	Status          v1.PodPhase
    29  	Namespace       string
    30  	Ans             map[string]string
    31  	Labels          map[string]string
    32  	ContainerStatus v1.ContainerStatus
    33  	Nodename        string
    34  }
    35  
    36  func NewPod(p PodArg) v1.Pod {
    37  	if p.Status == "" {
    38  		p.Status = v1.PodRunning
    39  	}
    40  	if p.Namespace == "" {
    41  		p.Namespace = metav1.NamespaceDefault
    42  	}
    43  	return v1.Pod{
    44  		TypeMeta: metav1.TypeMeta{
    45  			Kind:       "Pod",
    46  			APIVersion: "v1",
    47  		},
    48  		ObjectMeta: metav1.ObjectMeta{
    49  			Name:        p.Name,
    50  			Namespace:   p.Namespace,
    51  			Labels:      p.Labels,
    52  			Annotations: p.Ans,
    53  		},
    54  		Spec: v1.PodSpec{
    55  			NodeName: p.Nodename,
    56  		},
    57  		Status: v1.PodStatus{
    58  			Phase:             p.Status,
    59  			ContainerStatuses: []v1.ContainerStatus{p.ContainerStatus},
    60  		},
    61  	}
    62  }
    63  
    64  func GenerateNPods(
    65  	namePrefix string,
    66  	n int,
    67  	podArg PodArg,
    68  ) ([]runtime.Object, []v1.Pod) {
    69  	var podObjects []runtime.Object
    70  	var pods []v1.Pod
    71  	for i := 0; i < n; i++ {
    72  		podArg.Name = fmt.Sprintf("%s%d", namePrefix, i)
    73  		pod := NewPod(podArg)
    74  		podObjects = append(podObjects, &pod)
    75  		pods = append(pods, pod)
    76  	}
    77  
    78  	return podObjects, pods
    79  }
    80  
    81  func NewNode(
    82  	name string,
    83  	label map[string]string,
    84  ) v1.Node {
    85  	return v1.Node{
    86  		TypeMeta: metav1.TypeMeta{
    87  			Kind:       "Node",
    88  			APIVersion: "v1",
    89  		},
    90  		ObjectMeta: metav1.ObjectMeta{
    91  			Name:   name,
    92  			Labels: label,
    93  		},
    94  	}
    95  }
    96  
    97  func GenerateNNodes(
    98  	namePrefix string,
    99  	n int,
   100  	label map[string]string,
   101  ) ([]runtime.Object, []v1.Node) {
   102  	var nodeObjects []runtime.Object
   103  	var nodes []v1.Node
   104  
   105  	for i := 0; i < n; i++ {
   106  		node := NewNode(fmt.Sprintf("%s%d", namePrefix, i), label)
   107  		nodeObjects = append(nodeObjects, &node)
   108  		nodes = append(nodes, node)
   109  	}
   110  	return nodeObjects, nodes
   111  }
   112