generate.go

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

     1  // Copyright 2021 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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  
    16  package testutils
    17  
    18  import (
    19  	"fmt"
    20  
    21  	v1 "k8s.io/api/core/v1"
    22  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    23  	"k8s.io/apimachinery/pkg/runtime"
    24  
    25  	"github.com/chaos-mesh/chaos-mesh/api/v1alpha1"
    26  )
    27  
    28  // PodArg by default use `Status=corev1.PodRunning` and `Namespace=metav1.NamespaceDefault`.
    29  // For the others, the default values are empty.
    30  type PodArg struct {
    31  	Name            string
    32  	Status          v1.PodPhase
    33  	Namespace       string
    34  	Ans             map[string]string
    35  	Labels          map[string]string
    36  	ContainerStatus v1.ContainerStatus
    37  	Nodename        string
    38  }
    39  
    40  func NewPod(p PodArg) v1.Pod {
    41  	if p.Status == "" {
    42  		p.Status = v1.PodRunning
    43  	}
    44  	if p.Namespace == "" {
    45  		p.Namespace = metav1.NamespaceDefault
    46  	}
    47  	return v1.Pod{
    48  		TypeMeta: metav1.TypeMeta{
    49  			Kind:       "Pod",
    50  			APIVersion: "v1",
    51  		},
    52  		ObjectMeta: metav1.ObjectMeta{
    53  			Name:        p.Name,
    54  			Namespace:   p.Namespace,
    55  			Labels:      p.Labels,
    56  			Annotations: p.Ans,
    57  		},
    58  		Spec: v1.PodSpec{
    59  			NodeName: p.Nodename,
    60  		},
    61  		Status: v1.PodStatus{
    62  			Phase:             p.Status,
    63  			ContainerStatuses: []v1.ContainerStatus{p.ContainerStatus},
    64  		},
    65  	}
    66  }
    67  
    68  func GenerateNPods(
    69  	namePrefix string,
    70  	n int,
    71  	podArg PodArg,
    72  ) ([]runtime.Object, []v1.Pod) {
    73  	var podObjects []runtime.Object
    74  	var pods []v1.Pod
    75  	for i := 0; i < n; i++ {
    76  		podArg.Name = fmt.Sprintf("%s%d", namePrefix, i)
    77  		pod := NewPod(podArg)
    78  		podObjects = append(podObjects, &pod)
    79  		pods = append(pods, pod)
    80  	}
    81  
    82  	return podObjects, pods
    83  }
    84  
    85  func NewNode(
    86  	name string,
    87  	label map[string]string,
    88  ) v1.Node {
    89  	return v1.Node{
    90  		TypeMeta: metav1.TypeMeta{
    91  			Kind:       "Node",
    92  			APIVersion: "v1",
    93  		},
    94  		ObjectMeta: metav1.ObjectMeta{
    95  			Name:   name,
    96  			Labels: label,
    97  		},
    98  	}
    99  }
   100  
   101  func GenerateNNodes(
   102  	namePrefix string,
   103  	n int,
   104  	label map[string]string,
   105  ) ([]runtime.Object, []v1.Node) {
   106  	var nodeObjects []runtime.Object
   107  	var nodes []v1.Node
   108  
   109  	for i := 0; i < n; i++ {
   110  		node := NewNode(fmt.Sprintf("%s%d", namePrefix, i), label)
   111  		nodeObjects = append(nodeObjects, &node)
   112  		nodes = append(nodes, node)
   113  	}
   114  	return nodeObjects, nodes
   115  }
   116  
   117  // PhysicalMachineArg by default use `Namespace=metav1.NamespaceDefault`.
   118  // For the others, the default values are empty.
   119  type PhysicalMachineArg struct {
   120  	Name      string
   121  	Namespace string
   122  	Ans       map[string]string
   123  	Labels    map[string]string
   124  	Address   string
   125  }
   126  
   127  func NewPhysicalMachine(p PhysicalMachineArg) v1alpha1.PhysicalMachine {
   128  	if p.Namespace == "" {
   129  		p.Namespace = metav1.NamespaceDefault
   130  	}
   131  	return v1alpha1.PhysicalMachine{
   132  		TypeMeta: metav1.TypeMeta{
   133  			Kind:       "PhysicalMachine",
   134  			APIVersion: "v1alpha1",
   135  		},
   136  		ObjectMeta: metav1.ObjectMeta{
   137  			Name:        p.Name,
   138  			Namespace:   p.Namespace,
   139  			Labels:      p.Labels,
   140  			Annotations: p.Ans,
   141  		},
   142  		Spec: v1alpha1.PhysicalMachineSpec{
   143  			Address: p.Address,
   144  		},
   145  	}
   146  }
   147  
   148  func GenerateNPhysicalMachines(
   149  	namePrefix string,
   150  	n int,
   151  	arg PhysicalMachineArg,
   152  ) ([]runtime.Object, []v1alpha1.PhysicalMachine) {
   153  	var physicalMachineObjects []runtime.Object
   154  	var physicalMachines []v1alpha1.PhysicalMachine
   155  	for i := 0; i < n; i++ {
   156  		arg.Name = fmt.Sprintf("%s%d", namePrefix, i)
   157  		physicalMachine := NewPhysicalMachine(arg)
   158  		physicalMachineObjects = append(physicalMachineObjects, &physicalMachine)
   159  		physicalMachines = append(physicalMachines, physicalMachine)
   160  	}
   161  
   162  	return physicalMachineObjects, physicalMachines
   163  }
   164