feat(client): overwrite client context instead of setting new one (#20356)

Author: islishude

client/cmd.go

@@ -359,14 +359,17 @@ func GetClientContextFromCmd(cmd *cobra.Command) Context {
 // SetCmdClientContext sets a command's Context value to the provided argument.
 // If the context has not been set, set the given context as the default.
 func SetCmdClientContext(cmd *cobra.Command, clientCtx Context) error {
-	var cmdCtx context.Context
-
-	if cmd.Context() == nil {
+	cmdCtx := cmd.Context()
+	if cmdCtx == nil {
 		cmdCtx = context.Background()
+	}
+
+	v := cmd.Context().Value(ClientContextKey)
+	if clientCtxPtr, ok := v.(*Context); ok {
+		*clientCtxPtr = clientCtx
 	} else {
-		cmdCtx = cmd.Context()
+		cmd.SetContext(context.WithValue(cmdCtx, ClientContextKey, &clientCtx))
 	}
 
-	cmd.SetContext(context.WithValue(cmdCtx, ClientContextKey, &clientCtx))
 	return nil
 }

client/cmd_test.go

@@ -79,18 +79,21 @@ func TestSetCmdClientContextHandler(t *testing.T) {
 		name            string
 		expectedContext client.Context
 		args            []string
+		ctx             context.Context
 	}{
 		{
 			"no flags set",
 			initClientCtx,
 			[]string{},
+			context.WithValue(context.Background(), client.ClientContextKey, &client.Context{}),
 		},
 		{
 			"flags set",
 			initClientCtx.WithChainID("new-chain-id"),
 			[]string{
 				fmt.Sprintf("--%s=new-chain-id", flags.FlagChainID),
 			},
+			context.WithValue(context.Background(), client.ClientContextKey, &client.Context{}),
 		},
 		{
 			"flags set with space",
@@ -99,20 +102,37 @@ func TestSetCmdClientContextHandler(t *testing.T) {
 				fmt.Sprintf("--%s", flags.FlagHome),
 				"/tmp/dir",
 			},
+			context.Background(),
+		},
+		{
+			"no context provided",
+			initClientCtx.WithHomeDir("/tmp/noctx"),
+			[]string{
+				fmt.Sprintf("--%s", flags.FlagHome),
+				"/tmp/noctx",
+			},
+			nil,
+		},
+		{
+			"with invalid client value in the context",
+			initClientCtx.WithHomeDir("/tmp/invalid"),
+			[]string{
+				fmt.Sprintf("--%s", flags.FlagHome),
+				"/tmp/invalid",
+			},
+			context.WithValue(context.Background(), client.ClientContextKey, "invalid"),
 		},
 	}
 
 	for _, tc := range testCases {
 		tc := tc
 
 		t.Run(tc.name, func(t *testing.T) {
-			ctx := context.WithValue(context.Background(), client.ClientContextKey, &client.Context{})
-
 			cmd := newCmd()
 			_ = testutil.ApplyMockIODiscardOutErr(cmd)
 			cmd.SetArgs(tc.args)
 
-			require.NoError(t, cmd.ExecuteContext(ctx))
+			require.NoError(t, cmd.ExecuteContext(tc.ctx))
 
 			clientCtx := client.GetClientContextFromCmd(cmd)
 			require.Equal(t, tc.expectedContext, clientCtx)