split fund returns and keyfile update

cmd/seth.go

@@ -186,7 +186,7 @@ func RunCLI(args []string) error {
 						},
 						Action: func(cCtx *cli.Context) error {
 							toAddr := cCtx.String("address")
-							return seth.ReturnFunds(C, toAddr)
+							return seth.ReturnFundsAndUpdateKeyfile(C, toAddr)
 						},
 					},
 					{

keyfile.go

@@ -70,7 +70,7 @@ func UpdateAndSplitFunds(c *Client, opts *FundKeyFileCmdOpts) error {
 	return os.WriteFile(c.Cfg.KeyFilePath, b, os.ModePerm)
 }
 
-// ReturnFunds returns funds to the root key from all the test keys in some "keyfile.toml"
+// ReturnFundsAndUpdateKeyfile returns funds to the root key from all other keys
 func ReturnFunds(c *Client, toAddr string) error {
 	if toAddr == "" {
 		toAddr = c.Addresses[0].Hex()
@@ -140,14 +140,25 @@ func ReturnFunds(c *Client, toAddr string) error {
 	if err := eg.Wait(); err != nil {
 		return err
 	}
+
+	return nil
+}
+
+// ReturnFundsAndUpdateKeyfile returns funds to the root key from all the test keys in some "keyfile.toml"
+func ReturnFundsAndUpdateKeyfile(c *Client, toAddr string) error {
+	err := ReturnFunds(c, toAddr)
+	if err != nil {
+		return err
+	}
 	keyFile, err := c.CreateOrUnmarshalKeyFile(nil)
 	if err != nil {
 		return err
 	}
+	eg, egCtx := errgroup.WithContext(context.Background())
 	for _, kfd := range keyFile.Keys {
 		kfd := kfd
 		eg.Go(func() error {
-			balance, err := c.Client.BalanceAt(context.Background(), common.HexToAddress(kfd.Address), nil)
+			balance, err := c.Client.BalanceAt(egCtx, common.HexToAddress(kfd.Address), nil)
 			if err != nil {
 				return err
 			}