Skip to content

Commit 4bb7e7b

Browse files
hwdefhwdef
committed
add fakeclint in etcdctl, add UT for alarm command
Signed-off-by: hwdef <[email protected]>
1 parent 487e5b6 commit 4bb7e7b

File tree

3 files changed

+259
-2
lines changed

3 files changed

+259
-2
lines changed

‎etcdctl/ctlv3/command/alarm_command_test.go‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
// Copyright 2025 The etcd 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+
package command
16+
17+
import (
18+
"bytes"
19+
"context"
20+
"os"
21+
"strings"
22+
"testing"
23+
"time"
24+
25+
"github.com/spf13/cobra"
26+
27+
pb "go.etcd.io/etcd/api/v3/etcdserverpb"
28+
clientv3 "go.etcd.io/etcd/client/v3"
29+
"go.etcd.io/etcd/etcdctl/v3/ctlv3/command/fakeclient"
30+
)
31+
32+
func withStdoutCapture(t *testing.T, fn func()) string {
33+
t.Helper()
34+
old := os.Stdout
35+
r, w, _ := os.Pipe()
36+
os.Stdout = w
37+
defer func() { os.Stdout = old }()
38+
39+
fn()
40+
41+
w.Close()
42+
var buf bytes.Buffer
43+
_, _ = buf.ReadFrom(r)
44+
return buf.String()
45+
}
46+
47+
func newTestRoot() *cobra.Command {
48+
root := &cobra.Command{Use: "etcdctl"}
49+
root.AddGroup(NewKVGroup(), NewClusterMaintenanceGroup(), NewConcurrencyGroup(), NewAuthenticationGroup(), NewUtilityGroup())
50+
root.PersistentFlags().StringSlice("endpoints", []string{"127.0.0.1:2379"}, "")
51+
root.PersistentFlags().Bool("debug", false, "")
52+
root.PersistentFlags().String("write-out", "simple", "")
53+
root.PersistentFlags().Bool("hex", false, "")
54+
root.PersistentFlags().Duration("dial-timeout", time.Second, "")
55+
root.PersistentFlags().Duration("command-timeout", 5*time.Second, "")
56+
root.PersistentFlags().Duration("keepalive-time", time.Second, "")
57+
root.PersistentFlags().Duration("keepalive-timeout", 2*time.Second, "")
58+
root.PersistentFlags().Int("max-request-bytes", 0, "")
59+
root.PersistentFlags().Int("max-recv-bytes", 0, "")
60+
root.PersistentFlags().Bool("insecure-transport", true, "")
61+
root.PersistentFlags().Bool("insecure-discovery", true, "")
62+
root.PersistentFlags().Bool("insecure-skip-tls-verify", false, "")
63+
root.PersistentFlags().String("cert", "", "")
64+
root.PersistentFlags().String("key", "", "")
65+
root.PersistentFlags().String("cacert", "", "")
66+
root.PersistentFlags().String("auth-jwt-token", "", "")
67+
root.PersistentFlags().String("user", "", "")
68+
root.PersistentFlags().String("password", "", "")
69+
root.PersistentFlags().String("discovery-srv", "", "")
70+
root.PersistentFlags().String("discovery-srv-name", "", "")
71+
return root
72+
}
73+
74+
func TestAlarmList_PrintsAlarms(t *testing.T) {
75+
var gotReq int
76+
fc := &fakeclient.Client{
77+
AlarmListFn: func(ctx context.Context) (*clientv3.AlarmResponse, error) {
78+
gotReq++
79+
resp := clientv3.AlarmResponse{
80+
Alarms: []*pb.AlarmMember{
81+
{MemberID: 1, Alarm: pb.AlarmType_NOSPACE},
82+
{MemberID: 2, Alarm: pb.AlarmType_CORRUPT},
83+
},
84+
}
85+
return &resp, nil
86+
},
87+
}
88+
root := newTestRoot()
89+
root.SetContext(WithClient(context.Background(), fakeclient.WrapAsClientV3(fc)))
90+
root.AddCommand(NewAlarmCommand())
91+
root.SetArgs([]string{"alarm", "list"})
92+
93+
out := withStdoutCapture(t, func() { _ = root.Execute() })
94+
95+
if gotReq != 1 {
96+
t.Fatalf("expected AlarmList to be called once, got %d", gotReq)
97+
}
98+
if !strings.Contains(out, "NOSPACE") || !strings.Contains(out, "CORRUPT") {
99+
t.Fatalf("unexpected output: %q", out)
100+
}
101+
}
102+
103+
func TestAlarmDisarm_DisarmsAll(t *testing.T) {
104+
var disarmCalled int
105+
fc := &fakeclient.Client{
106+
AlarmDisarmFn: func(ctx context.Context, m *clientv3.AlarmMember) (*clientv3.AlarmResponse, error) {
107+
disarmCalled++
108+
return &clientv3.AlarmResponse{Alarms: []*pb.AlarmMember{}}, nil
109+
},
110+
}
111+
root := newTestRoot()
112+
root.SetContext(WithClient(context.Background(), fakeclient.WrapAsClientV3(fc)))
113+
root.AddCommand(NewAlarmCommand())
114+
root.SetArgs([]string{"alarm", "disarm"})
115+
116+
_ = withStdoutCapture(t, func() { _ = root.Execute() })
117+
118+
if disarmCalled != 1 {
119+
t.Fatalf("expected disarm to be called, got %d", disarmCalled)
120+
}
121+
}

‎etcdctl/ctlv3/command/fakeclient/client.go‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
// Copyright 2025 The etcd 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+
package fakeclient
16+
17+
import (
18+
"bytes"
19+
"context"
20+
"errors"
21+
"io"
22+
23+
clientv3 "go.etcd.io/etcd/client/v3"
24+
)
25+
26+
// Client is a lightweight fake that satisfies the small subset of etcdctl usage.
27+
// It embeds function fields for behaviors; unimplemented methods return error.
28+
// This keeps the fake reusable across most commands by stubbing only used calls.
29+
type Client struct {
30+
AlarmListFn func(ctx context.Context) (*clientv3.AlarmResponse, error)
31+
AlarmDisarmFn func(ctx context.Context, m *clientv3.AlarmMember) (*clientv3.AlarmResponse, error)
32+
}
33+
34+
// New returns a fake client instance. Config is ignored.
35+
func New(_ clientv3.Config) (*clientv3.Client, error) {
36+
return nil, errors.New("fakeclient.New must be used through command-level wrapper in tests")
37+
}
38+
39+
// WrapAsClientV3 constructs a minimal *clientv3.Client using NewCtxClient and attaches fake Maintenance.
40+
func WrapAsClientV3(fc *Client) *clientv3.Client {
41+
c := clientv3.NewCtxClient(context.Background())
42+
c.Maintenance = &maintenance{fake: fc}
43+
return c
44+
}
45+
46+
type maintenance struct{ fake *Client }
47+
48+
func (m *maintenance) AlarmList(ctx context.Context) (*clientv3.AlarmResponse, error) {
49+
if m.fake.AlarmListFn != nil {
50+
return m.fake.AlarmListFn(ctx)
51+
}
52+
return &clientv3.AlarmResponse{}, nil
53+
}
54+
55+
func (m *maintenance) AlarmDisarm(ctx context.Context, am *clientv3.AlarmMember) (*clientv3.AlarmResponse, error) {
56+
if m.fake.AlarmDisarmFn != nil {
57+
return m.fake.AlarmDisarmFn(ctx, am)
58+
}
59+
return &clientv3.AlarmResponse{}, nil
60+
}
61+
62+
func (m *maintenance) Defragment(ctx context.Context, endpoint string) (*clientv3.DefragmentResponse, error) {
63+
var resp clientv3.DefragmentResponse
64+
return &resp, nil
65+
}
66+
67+
func (m *maintenance) Status(ctx context.Context, endpoint string) (*clientv3.StatusResponse, error) {
68+
var resp clientv3.StatusResponse
69+
return &resp, nil
70+
}
71+
72+
func (m *maintenance) HashKV(ctx context.Context, endpoint string, rev int64) (*clientv3.HashKVResponse, error) {
73+
var resp clientv3.HashKVResponse
74+
return &resp, nil
75+
}
76+
77+
func (m *maintenance) SnapshotWithVersion(ctx context.Context) (*clientv3.SnapshotResponse, error) {
78+
return &clientv3.SnapshotResponse{Snapshot: io.NopCloser(bytes.NewReader(nil))}, nil
79+
}
80+
81+
func (m *maintenance) Snapshot(ctx context.Context) (io.ReadCloser, error) {
82+
return io.NopCloser(bytes.NewReader(nil)), nil
83+
}
84+
85+
func (m *maintenance) MoveLeader(ctx context.Context, transfereeID uint64) (*clientv3.MoveLeaderResponse, error) {
86+
var resp clientv3.MoveLeaderResponse
87+
return &resp, nil
88+
}
89+
90+
func (m *maintenance) Downgrade(ctx context.Context, action clientv3.DowngradeAction, version string) (*clientv3.DowngradeResponse, error) {
91+
var resp clientv3.DowngradeResponse
92+
return &resp, nil
93+
}

‎etcdctl/ctlv3/command/global.go‎

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
package command
1616

1717
import (
18+
"context"
1819
"errors"
1920
"fmt"
2021
"io"
@@ -71,6 +72,44 @@ type discoveryCfg struct {
7172

7273
var display printer = &simplePrinter{}
7374

75+
var newClientFunc = clientv3.New
76+
77+
type ClientFactory func(clientv3.Config) (*clientv3.Client, error)
78+
79+
type clientFactoryKey struct{}
80+
81+
// WithClientFactory attaches a custom client factory to the provided context.
82+
// Tests can inject fakes without mutating the global client constructor.
83+
func WithClientFactory(ctx context.Context, factory ClientFactory) context.Context {
84+
if ctx == nil {
85+
ctx = context.Background()
86+
}
87+
if factory == nil {
88+
return ctx
89+
}
90+
return context.WithValue(ctx, clientFactoryKey{}, factory)
91+
}
92+
93+
func WithClient(ctx context.Context, cli *clientv3.Client) context.Context {
94+
if cli == nil {
95+
return ctx
96+
}
97+
return WithClientFactory(ctx, func(clientv3.Config) (*clientv3.Client, error) {
98+
return cli, nil
99+
})
100+
}
101+
102+
func clientFactoryFromCmd(cmd *cobra.Command) ClientFactory {
103+
if cmd != nil {
104+
if ctx := cmd.Context(); ctx != nil {
105+
if factory, ok := ctx.Value(clientFactoryKey{}).(ClientFactory); ok && factory != nil {
106+
return factory
107+
}
108+
}
109+
}
110+
return newClientFunc
111+
}
112+
74113
func initDisplayFromCmd(cmd *cobra.Command) {
75114
isHex, err := cmd.Flags().GetBool("hex")
76115
if err != nil {
@@ -153,17 +192,21 @@ func mustClientCfgFromCmd(cmd *cobra.Command) *clientv3.Config {
153192

154193
func mustClientFromCmd(cmd *cobra.Command) *clientv3.Client {
155194
cfg := clientConfigFromCmd(cmd)
156-
return mustClient(cfg)
195+
return mustClientWithFactory(cmd, cfg)
157196
}
158197

159198
func mustClient(cc *clientv3.ConfigSpec) *clientv3.Client {
199+
return mustClientWithFactory(nil, cc)
200+
}
201+
202+
func mustClientWithFactory(cmd *cobra.Command, cc *clientv3.ConfigSpec) *clientv3.Client {
160203
lg, _ := logutil.CreateDefaultZapLogger(zap.InfoLevel)
161204
cfg, err := clientv3.NewClientConfig(cc, lg)
162205
if err != nil {
163206
cobrautl.ExitWithError(cobrautl.ExitBadArgs, err)
164207
}
165208

166-
client, err := clientv3.New(*cfg)
209+
client, err := clientFactoryFromCmd(cmd)(*cfg)
167210
if err != nil {
168211
cobrautl.ExitWithError(cobrautl.ExitBadConnection, err)
169212
}

0 commit comments

Comments
 (0)