chore: improve pkg kbagent coverage (#10386)

Author: leon-ape

pkg/kbagent/client/client_test.go

@@ -0,0 +1,166 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package client
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/go-logr/logr"
+	"github.com/golang/mock/gomock"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/client-go/rest"
+
+	"github.com/apecloud/kubeblocks/pkg/kbagent/proto"
+)
+
+type stubClient struct{}
+
+func (stubClient) Close() error {
+	return nil
+}
+
+func (stubClient) Action(context.Context, proto.ActionRequest) (proto.ActionResponse, error) {
+	return proto.ActionResponse{Message: "ok"}, nil
+}
+
+func TestMockClientLifecycle(t *testing.T) {
+	t.Cleanup(UnsetMockClient)
+
+	mock := stubClient{}
+	SetMockClient(mock, nil)
+	if GetMockClient() != mock {
+		t.Fatalf("GetMockClient() did not return mock client")
+	}
+	got, err := NewClient(func() (string, int32, error) {
+		t.Fatal("endpoint should not be called when mock client is set")
+		return "", 0, nil
+	})
+	if err != nil || got != mock {
+		t.Fatalf("NewClient() = %v, %v, want mock nil-error", got, err)
+	}
+
+	mockErr := errors.New("mock")
+	SetMockClient(nil, mockErr)
+	got, err = NewClient(func() (string, int32, error) {
+		t.Fatal("endpoint should not be called when mock error is set")
+		return "", 0, nil
+	})
+	if got != nil || !errors.Is(err, mockErr) {
+		t.Fatalf("NewClient() = %v, %v, want nil mockErr", got, err)
+	}
+
+	UnsetMockClient()
+	if GetMockClient() != nil {
+		t.Fatalf("mock client should be cleared")
+	}
+}
+
+func TestNewClientEndpointBranches(t *testing.T) {
+	endpointErr := errors.New("endpoint")
+	if got, err := NewClient(func() (string, int32, error) { return "", 0, endpointErr }); got != nil || !errors.Is(err, endpointErr) {
+		t.Fatalf("NewClient endpoint error = %v, %v", got, err)
+	}
+
+	if got, err := NewClient(func() (string, int32, error) { return "", 0, nil }); got != nil || err != nil {
+		t.Fatalf("NewClient empty endpoint = %v, %v, want nil nil", got, err)
+	}
+
+	got, err := NewClient(func() (string, int32, error) { return "127.0.0.1", 3501, nil })
+	if err != nil {
+		t.Fatalf("NewClient returned error: %v", err)
+	}
+	if _, ok := got.(*httpClient); !ok {
+		t.Fatalf("NewClient returned %T, want *httpClient", got)
+	}
+	if err := got.Close(); err != nil {
+		t.Fatalf("Close() error = %v", err)
+	}
+}
+
+func TestGeneratedMockClient(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mock := NewMockClient(ctrl)
+	mock.EXPECT().
+		Action(gomock.Any(), proto.ActionRequest{Action: "backup"}).
+		Return(proto.ActionResponse{Message: "ok"}, nil)
+
+	resp, err := mock.Action(context.Background(), proto.ActionRequest{Action: "backup"})
+	if err != nil {
+		t.Fatalf("mock Action() error = %v", err)
+	}
+	if resp.Message != "ok" {
+		t.Fatalf("mock Action() response = %#v", resp)
+	}
+	if err := mock.Close(); err != nil {
+		t.Fatalf("mock Close() error = %v", err)
+	}
+}
+
+func TestNewPortForwardClientStableBranches(t *testing.T) {
+	t.Cleanup(UnsetMockClient)
+
+	mock := stubClient{}
+	SetMockClient(mock, nil)
+	got, err := NewPortForwardClient(&corev1.Pod{}, func() (string, int32, error) {
+		t.Fatal("endpoint should not be called when mock client is set")
+		return "", 0, nil
+	})
+	if err != nil || got != mock {
+		t.Fatalf("NewPortForwardClient mock = %v, %v", got, err)
+	}
+
+	UnsetMockClient()
+	endpointErr := errors.New("endpoint")
+	got, err = NewPortForwardClient(&corev1.Pod{}, func() (string, int32, error) {
+		return "", 0, endpointErr
+	})
+	if got != nil || !errors.Is(err, endpointErr) {
+		t.Fatalf("NewPortForwardClient endpoint error = %v, %v", got, err)
+	}
+}
+
+func TestPortForwardClientStableErrors(t *testing.T) {
+	pf := &portForwardClient{
+		pod:    &corev1.Pod{},
+		port:   "3501",
+		config: &rest.Config{},
+		logger: logr.Discard(),
+	}
+	if err := pf.Close(); err != nil {
+		t.Fatalf("Close() error = %v", err)
+	}
+	_, _ = pf.createDialer("POST", nil, &rest.Config{})
+	readyCh := make(chan struct{})
+	stopCh := make(chan struct{})
+	forwarder, err := pf.newPortForwarder(readyCh, stopCh, io.Discard)
+	close(stopCh)
+	if err != nil {
+		t.Fatalf("newPortForwarder() error = %v", err)
+	}
+	if forwarder == nil {
+		t.Fatalf("newPortForwarder() returned nil")
+	}
+	if resp, err := pf.Action(context.Background(), proto.ActionRequest{Action: "backup"}); err == nil || resp.Message != "" {
+		t.Fatalf("expected Action error for empty rest config, got %#v, %v", resp, err)
+	}
+}

pkg/kbagent/client/http_client_test.go

@@ -0,0 +1,128 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package client
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"testing"
+
+	"github.com/apecloud/kubeblocks/pkg/constant"
+	"github.com/apecloud/kubeblocks/pkg/kbagent/proto"
+)
+
+type errorReader struct{}
+
+func (errorReader) Read([]byte) (int, error) {
+	return 0, errors.New("read")
+}
+
+func newHTTPClientForTest(t *testing.T, handler http.HandlerFunc) (*httpClient, func()) {
+	t.Helper()
+	server := httptest.NewServer(handler)
+	u, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatalf("parse server URL: %v", err)
+	}
+	host, portString, err := net.SplitHostPort(u.Host)
+	if err != nil {
+		t.Fatalf("split host port: %v", err)
+	}
+	var port int32
+	if _, err := fmt.Sscan(portString, &port); err != nil {
+		t.Fatalf("parse port: %v", err)
+	}
+	return &httpClient{host: host, port: port, client: server.Client()}, server.Close
+}
+
+func TestHTTPClientAction(t *testing.T) {
+	cli, closeServer := newHTTPClientForTest(t, func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != proto.ServiceAction.URI || r.Method != http.MethodPost {
+			t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
+		}
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"message":"done","output":"b2s="}`))
+	})
+	defer closeServer()
+
+	resp, err := cli.Action(context.Background(), proto.ActionRequest{Action: "backup"})
+	if err != nil {
+		t.Fatalf("Action() error = %v", err)
+	}
+	if resp.Message != "done" || string(resp.Output) != "ok" {
+		t.Fatalf("unexpected response: %#v", resp)
+	}
+
+	resp, err = cli.Action(context.WithValue(context.Background(), constant.DryRunContextKey, true), proto.ActionRequest{Action: "backup"})
+	if err != nil || resp.Message != "" || resp.Error != "" || len(resp.Output) != 0 {
+		t.Fatalf("dry-run Action() = %#v, %v", resp, err)
+	}
+}
+
+func TestHTTPClientRequestAndDecodeErrors(t *testing.T) {
+	cli, closeServer := newHTTPClientForTest(t, func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Query().Get("case") {
+		case "bad-status":
+			w.WriteHeader(http.StatusAccepted)
+		case "invalid-json":
+			w.WriteHeader(http.StatusInternalServerError)
+			_, _ = w.Write([]byte("{"))
+		default:
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte(`{"error":"failed"}`))
+		}
+	})
+	defer closeServer()
+
+	if _, err := cli.request(context.Background(), " bad method", "http://bad-url", nil); err == nil {
+		t.Fatalf("expected request construction error")
+	}
+
+	body, err := cli.request(context.Background(), http.MethodPost, "http://"+net.JoinHostPort(cli.host, fmt.Sprint(cli.port))+proto.ServiceAction.URI, strings.NewReader("{}"))
+	if err != nil {
+		t.Fatalf("request() error = %v", err)
+	}
+	_ = body.Close()
+
+	_, err = cli.request(context.Background(), http.MethodPost, "http://"+net.JoinHostPort(cli.host, fmt.Sprint(cli.port))+proto.ServiceAction.URI+"?case=bad-status", nil)
+	if err == nil {
+		t.Fatalf("expected unexpected status error")
+	}
+
+	body, err = cli.request(context.Background(), http.MethodPost, "http://"+net.JoinHostPort(cli.host, fmt.Sprint(cli.port))+proto.ServiceAction.URI+"?case=invalid-json", nil)
+	if err != nil {
+		t.Fatalf("request invalid-json error = %v", err)
+	}
+	defer body.Close()
+	if _, err := decode(body, &proto.ActionResponse{}); err == nil {
+		t.Fatalf("expected decode error")
+	}
+
+	if _, err := decode(io.NopCloser(errorReader{}), &proto.ActionResponse{}); err == nil {
+		t.Fatalf("expected read error")
+	}
+}

pkg/kbagent/proto/errors_test.go

@@ -0,0 +1,79 @@
+/*
+Copyright (C) 2022-2026 ApeCloud Co., Ltd
+
+This file is part of KubeBlocks project
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+package proto
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestError2Type(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want string
+	}{
+		{name: "nil", want: ""},
+		{name: "not defined", err: ErrNotDefined, want: "notDefined"},
+		{name: "not implemented", err: ErrNotImplemented, want: "notImplemented"},
+		{name: "precondition failed", err: ErrPreconditionFailed, want: "preconditionFailed"},
+		{name: "bad request", err: ErrBadRequest, want: "badRequest"},
+		{name: "in progress", err: ErrInProgress, want: "inProgress"},
+		{name: "busy", err: ErrBusy, want: "busy"},
+		{name: "timed out", err: ErrTimedOut, want: "timedOut"},
+		{name: "failed", err: ErrFailed, want: "failed"},
+		{name: "internal error", err: ErrInternalError, want: "internalError"},
+		{name: "wrapped", err: errors.Join(ErrBusy), want: "busy"},
+		{name: "unknown", err: errors.New("other"), want: "unknown"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := Error2Type(tt.err); got != tt.want {
+				t.Fatalf("Error2Type() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestType2Error(t *testing.T) {
+	tests := []struct {
+		errType string
+		want    error
+	}{
+		{errType: "", want: nil},
+		{errType: "notDefined", want: ErrNotDefined},
+		{errType: "notImplemented", want: ErrNotImplemented},
+		{errType: "preconditionFailed", want: ErrPreconditionFailed},
+		{errType: "badRequest", want: ErrBadRequest},
+		{errType: "inProgress", want: ErrInProgress},
+		{errType: "busy", want: ErrBusy},
+		{errType: "timedOut", want: ErrTimedOut},
+		{errType: "failed", want: ErrFailed},
+		{errType: "internalError", want: ErrInternalError},
+		{errType: "unknown-type", want: ErrUnknown},
+	}
+	for _, tt := range tests {
+		t.Run(tt.errType, func(t *testing.T) {
+			if got := Type2Error(tt.errType); !errors.Is(got, tt.want) {
+				t.Fatalf("Type2Error() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}