github.com/k8snetworkplumbingwg/sriov-network-operator@v1.2.1-0.20240408194816-2d2e5a45d453/test/util/pod/pod.go (about)

     1  package pod
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"strings"
     8  	"time"
     9  
    10  	corev1 "k8s.io/api/core/v1"
    11  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    12  	"k8s.io/client-go/kubernetes/scheme"
    13  	"k8s.io/client-go/tools/remotecommand"
    14  	"k8s.io/utils/pointer"
    15  
    16  	testclient "github.com/k8snetworkplumbingwg/sriov-network-operator/test/util/client"
    17  	"github.com/k8snetworkplumbingwg/sriov-network-operator/test/util/images"
    18  	"github.com/k8snetworkplumbingwg/sriov-network-operator/test/util/namespaces"
    19  )
    20  
    21  const hostnameLabel = "kubernetes.io/hostname"
    22  
    23  func GetDefinition() *corev1.Pod {
    24  	podObject := &corev1.Pod{
    25  		ObjectMeta: metav1.ObjectMeta{
    26  			GenerateName: "testpod-",
    27  			Namespace:    namespaces.Test},
    28  		Spec: corev1.PodSpec{
    29  			TerminationGracePeriodSeconds: pointer.Int64Ptr(0),
    30  			Containers: []corev1.Container{{Name: "test",
    31  				Image: images.Test(),
    32  				SecurityContext: &corev1.SecurityContext{
    33  					Capabilities: &corev1.Capabilities{
    34  						Add: []corev1.Capability{"NET_RAW"},
    35  					}},
    36  				Command: []string{"/bin/bash", "-c", "sleep INF"}}}}}
    37  
    38  	return podObject
    39  }
    40  
    41  func DefineWithNetworks(networks []string) *corev1.Pod {
    42  	podObject := GetDefinition()
    43  	podObject.Annotations = map[string]string{"k8s.v1.cni.cncf.io/networks": strings.Join(networks, ",")}
    44  
    45  	return podObject
    46  }
    47  
    48  func DefineWithHostNetwork(nodeName string) *corev1.Pod {
    49  	podObject := GetDefinition()
    50  	podObject.Spec.HostNetwork = true
    51  	podObject.Spec.NodeSelector = map[string]string{
    52  		"kubernetes.io/hostname": nodeName,
    53  	}
    54  
    55  	return podObject
    56  }
    57  
    58  // RedefineAsPrivileged updates the pod to be privileged
    59  func RedefineAsPrivileged(pod *corev1.Pod) *corev1.Pod {
    60  	pod.Spec.Containers[0].SecurityContext = &corev1.SecurityContext{}
    61  	b := true
    62  	pod.Spec.Containers[0].SecurityContext.Privileged = &b
    63  	return pod
    64  }
    65  
    66  // RedefineWithHostNetwork updates the pod definition Spec.HostNetwork to true
    67  func RedefineWithHostNetwork(pod *corev1.Pod) *corev1.Pod {
    68  	pod.Spec.HostNetwork = true
    69  	return pod
    70  }
    71  
    72  // RedefineWithNodeSelector updates the pod definition with a node selector
    73  func RedefineWithNodeSelector(pod *corev1.Pod, node string) *corev1.Pod {
    74  	pod.Spec.NodeSelector = map[string]string{
    75  		hostnameLabel: node,
    76  	}
    77  	return pod
    78  }
    79  
    80  // RedefineWithMount updates the pod definition with a volume and volume mount
    81  func RedefineWithMount(pod *corev1.Pod, volume corev1.Volume, mount corev1.VolumeMount) *corev1.Pod {
    82  	pod.Spec.Containers[0].VolumeMounts = []corev1.VolumeMount{mount}
    83  	pod.Spec.Volumes = []corev1.Volume{volume}
    84  
    85  	return pod
    86  }
    87  
    88  // RedefineWithCommand updates the pod definition with a different command
    89  func RedefineWithCommand(pod *corev1.Pod, command []string, args []string) *corev1.Pod {
    90  	pod.Spec.Containers[0].Command = command
    91  	pod.Spec.Containers[0].Args = args
    92  	return pod
    93  }
    94  
    95  // RedefineWithRestartPolicy updates the pod definition with a restart policy
    96  func RedefineWithRestartPolicy(pod *corev1.Pod, restartPolicy corev1.RestartPolicy) *corev1.Pod {
    97  	pod.Spec.RestartPolicy = restartPolicy
    98  	return pod
    99  }
   100  
   101  func RedefineWithCapabilities(pod *corev1.Pod, capabilitiesList []corev1.Capability) *corev1.Pod {
   102  	pod.Spec.Containers[0].SecurityContext = &corev1.SecurityContext{Capabilities: &corev1.Capabilities{Add: capabilitiesList}}
   103  	return pod
   104  }
   105  
   106  // ExecCommand runs command in the pod and returns buffer output
   107  func ExecCommand(cs *testclient.ClientSet, pod *corev1.Pod, command ...string) (string, string, error) {
   108  	var buf, errbuf bytes.Buffer
   109  	req := cs.CoreV1Interface.RESTClient().
   110  		Post().
   111  		Namespace(pod.Namespace).
   112  		Resource("pods").
   113  		Name(pod.Name).
   114  		SubResource("exec").
   115  		VersionedParams(&corev1.PodExecOptions{
   116  			Container: pod.Spec.Containers[0].Name,
   117  			Command:   command,
   118  			Stdout:    true,
   119  			Stderr:    true,
   120  			TTY:       true,
   121  		}, scheme.ParameterCodec)
   122  
   123  	exec, err := remotecommand.NewSPDYExecutor(cs.Config, "POST", req.URL())
   124  	if err != nil {
   125  		return buf.String(), errbuf.String(), err
   126  	}
   127  
   128  	err = exec.Stream(remotecommand.StreamOptions{
   129  		Stdout: &buf,
   130  		Stderr: &errbuf,
   131  		Tty:    true,
   132  	})
   133  	if err != nil {
   134  		return buf.String(), errbuf.String(), err
   135  	}
   136  
   137  	return buf.String(), errbuf.String(), nil
   138  }
   139  
   140  // GetLog connects to a pod and fetches log
   141  func GetLog(cs *testclient.ClientSet, p *corev1.Pod, s time.Duration) (string, error) {
   142  	logStart := int64(s.Seconds())
   143  	req := cs.Pods(p.Namespace).GetLogs(p.Name, &corev1.PodLogOptions{SinceSeconds: &logStart})
   144  	log, err := req.Stream(context.Background())
   145  	if err != nil {
   146  		return "", err
   147  	}
   148  	defer log.Close()
   149  
   150  	buf := new(bytes.Buffer)
   151  	_, err = io.Copy(buf, log)
   152  
   153  	if err != nil {
   154  		return "", err
   155  	}
   156  
   157  	return buf.String(), nil
   158  }