Add test for kubeLoggingWithReconnect

Author: matoval

go.mod

@@ -62,6 +62,7 @@ require (
 	github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f // indirect
 	github.com/onsi/ginkgo/v2 v2.21.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
 	github.com/sourcegraph/conc v0.3.0 // indirect
@@ -83,6 +84,7 @@ require (
 	golang.org/x/time v0.5.0 // indirect
 	golang.org/x/tools v0.27.0 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
+	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect

pkg/workceptor/kubernetes.go

@@ -294,7 +294,7 @@ func (kw *KubeUnit) kubeLoggingNoReconnect(streamWait *sync.WaitGroup, stdout *S
 	}
 }
 
-func (kw *KubeUnit) kubeLoggingWithReconnect(streamWait *sync.WaitGroup, stdout *STDoutWriter, stdinErr *error, stdoutErr *error) {
+func (kw *KubeUnit) KubeLoggingWithReconnect(streamWait *sync.WaitGroup, stdout *STDoutWriter, stdinErr *error, stdoutErr *error) {
 	// preferred method for k8s >= 1.23.14
 	defer streamWait.Done()
 	var sinceTime time.Time
@@ -420,7 +420,7 @@ func (kw *KubeUnit) kubeLoggingWithReconnect(streamWait *sync.WaitGroup, stdout
 	}
 }
 
-func (kw *KubeUnit) createPod(env map[string]string) error {
+func (kw *KubeUnit) CreatePod(env map[string]string) error {
 	ked := kw.UnredactedStatus().ExtraData.(*KubeExtraData)
 	command, err := shlex.Split(ked.Command)
 	if err != nil {
@@ -621,7 +621,7 @@ func (kw *KubeUnit) runWorkUsingLogger() {
 	if podName == "" {
 		// create new pod if ked.PodName is empty
 		// TODO: add retry logic to make this more resilient to transient errors
-		if err := kw.createPod(nil); err != nil {
+		if err := kw.CreatePod(nil); err != nil {
 			if err != ErrPodCompleted {
 				errMsg := fmt.Sprintf("Error creating pod: %s", err)
 				kw.GetWorkceptor().nc.GetLogger().Error(errMsg) //nolint:govet
@@ -842,7 +842,7 @@ func (kw *KubeUnit) runWorkUsingLogger() {
 	stdoutWithReconnect := ShouldUseReconnect(kw)
 	if stdoutWithReconnect && stdoutErr == nil {
 		kw.GetWorkceptor().nc.GetLogger().Debug("streaming stdout with reconnect support")
-		go kw.kubeLoggingWithReconnect(&streamWait, stdout, &stdinErr, &stdoutErr)
+		go kw.KubeLoggingWithReconnect(&streamWait, stdout, &stdinErr, &stdoutErr)
 	} else {
 		kw.GetWorkceptor().nc.GetLogger().Debug("streaming stdout with no reconnect support")
 		go kw.kubeLoggingNoReconnect(&streamWait, stdout, &stdoutErr)
@@ -1062,7 +1062,7 @@ func (kw *KubeUnit) runWorkUsingTCP() {
 	}()
 
 	// Create the pod
-	err = kw.createPod(map[string]string{"RECEPTOR_HOST": listenHost, "RECEPTOR_PORT": listenPort})
+	err = kw.CreatePod(map[string]string{"RECEPTOR_HOST": listenHost, "RECEPTOR_PORT": listenPort})
 	if err != nil {
 		errMsg := fmt.Sprintf("Error creating pod: %s", err)
 		kw.UpdateBasicStatus(WorkStateFailed, errMsg, 0)

pkg/workceptor/kubernetes_test.go

@@ -20,6 +20,7 @@ import (
 	"k8s.io/apimachinery/pkg/selection"
 	"k8s.io/apimachinery/pkg/watch"
 	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/kubernetes/fake"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/remotecommand"
 )
@@ -423,3 +424,64 @@ func Test_IsCompatibleK8S(t *testing.T) {
 		})
 	}
 }
+
+func TestKubeLoggingWithReconnect(t *testing.T) {
+	var stdinErr error
+	var stdoutErr error
+	ctx := context.Background()
+	ctrl := gomock.NewController(t)
+	mockBaseWorkUnitForWorkUnit := mock_workceptor.NewMockBaseWorkUnitForWorkUnit(ctrl)
+	mockBaseWorkUnit := mock_workceptor.NewMockBaseWorkUnitForWorkUnit(ctrl)
+	mockNetceptor := mock_workceptor.NewMockNetceptorForWorkceptor(ctrl)
+	mockNetceptor.EXPECT().NodeID().Return("NodeID")
+	mockKubeAPI := mock_workceptor.NewMockKubeAPIer(ctrl)
+
+	kw := &workceptor.KubeUnit{
+		BaseWorkUnitForWorkUnit: mockBaseWorkUnitForWorkUnit,
+	}
+
+	w, err := workceptor.New(ctx, mockNetceptor, "/tmp")
+	if err != nil {
+		t.Errorf("Error while creating Workceptor: %v", err)
+	}
+
+	mockBaseWorkUnit.EXPECT().Init(w, "", "", workceptor.FileSystem{}, nil)
+	kubeConfig := workceptor.KubeWorkerCfg{AuthMethod: "incluster"}
+	kubeConfig.NewkubeWorker(mockBaseWorkUnit, w, "", "", mockKubeAPI)
+
+	tests := []struct {
+		name     string
+		expectedCalls func()
+	}{
+		{
+			name:     "Kube error should be read",
+			expectedCalls: func() {
+				lock := &sync.RWMutex{}
+				mockBaseWorkUnitForWorkUnit.EXPECT().GetStatusLock().Return(lock).AnyTimes()
+				kubeExtraData := workceptor.KubeExtraData{}
+				status := workceptor.StatusFileData{ExtraData: &kubeExtraData}
+				mockBaseWorkUnitForWorkUnit.EXPECT().GetStatusWithoutExtraData().Return(&status).AnyTimes()
+				mockBaseWorkUnitForWorkUnit.EXPECT().GetStatusCopy().Return(status).AnyTimes()
+				mockBaseWorkUnitForWorkUnit.EXPECT().GetContext().Return(ctx).AnyTimes()
+				pod := corev1.Pod{TypeMeta: metav1.TypeMeta{}, ObjectMeta: metav1.ObjectMeta{Name: "Test Name"}, Spec: corev1.PodSpec{}, Status: corev1.PodStatus{}}
+				mockKubeAPI.EXPECT().Create(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&pod, nil).AnyTimes()
+				mockBaseWorkUnitForWorkUnit.EXPECT().UpdateFullStatus(gomock.Any()).AnyTimes()
+				field := hasTerm{}
+				mockKubeAPI.EXPECT().OneTermEqualSelector(gomock.Any(), gomock.Any()).Return(&field).AnyTimes()
+				ev := watch.Event{Object: &pod}
+				mockKubeAPI.EXPECT().UntilWithSync(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&ev, nil).AnyTimes()
+				mockKubeAPI.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&pod, nil).AnyTimes()
+				client := fake.NewSimpleClientset()
+				req := client.CoreV1().Pods("default").GetLogs("pod-1", &corev1.PodLogOptions{})
+				mockKubeAPI.EXPECT().GetLogs(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(req)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.expectedCalls()
+			kw.CreatePod(nil)
+			kw.KubeLoggingWithReconnect(&sync.WaitGroup{}, &workceptor.STDoutWriter{}, &stdinErr, &stdoutErr)
+		})
+	}
+}