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

This commit is contained in:
Shude Li 2024-05-14 21:27:59 +08:00 committed by GitHub
parent 7ae23e287a
commit b2fe0ffec6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 32 additions and 9 deletions

View File

@ -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()
} else {
cmdCtx = cmd.Context()
}
cmd.SetContext(context.WithValue(cmdCtx, ClientContextKey, &clientCtx))
v := cmd.Context().Value(ClientContextKey)
if clientCtxPtr, ok := v.(*Context); ok {
*clientCtxPtr = clientCtx
} else {
cmd.SetContext(context.WithValue(cmdCtx, ClientContextKey, &clientCtx))
}
return nil
}

View File

@ -79,11 +79,13 @@ 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",
@ -91,6 +93,7 @@ func TestSetCmdClientContextHandler(t *testing.T) {
[]string{
fmt.Sprintf("--%s=new-chain-id", flags.FlagChainID),
},
context.WithValue(context.Background(), client.ClientContextKey, &client.Context{}),
},
{
"flags set with space",
@ -99,6 +102,25 @@ 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"),
},
}
@ -106,13 +128,11 @@ func TestSetCmdClientContextHandler(t *testing.T) {
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)