ethclient/gethclient: fix bugs in override object encoding

ethclient/gethclient/gethclient.go

@@ -216,16 +216,21 @@ func toOverrideMap(overrides *map[common.Address]OverrideAccount) interface{} {
 	}
 	type overrideAccount struct {
 		Nonce     hexutil.Uint64              `json:"nonce"`
-		Code      hexutil.Bytes               `json:"code"`
+		Code      string                      `json:"code,omitempty"`
 		Balance   *hexutil.Big                `json:"balance"`
 		State     map[common.Hash]common.Hash `json:"state"`
 		StateDiff map[common.Hash]common.Hash `json:"stateDiff"`
 	}
 	result := make(map[common.Address]overrideAccount)
+
 	for addr, override := range *overrides {
+		var code string
+		if override.Code != nil {
+			code = hexutil.Encode(override.Code)
+		}
 		result[addr] = overrideAccount{
 			Nonce:     hexutil.Uint64(override.Nonce),
-			Code:      override.Code,
+			Code:      code,
 			Balance:   (*hexutil.Big)(override.Balance),
 			State:     override.State,
 			StateDiff: override.StateDiff,

ethclient/gethclient/gethclient_test.go

@@ -19,11 +19,13 @@ package gethclient
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"math/big"
 	"testing"
 
 	"github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/consensus/ethash"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/rawdb"
@@ -135,6 +137,12 @@ func TestGethClient(t *testing.T) {
 		}, {
 			"TestCallContract",
 			func(t *testing.T) { testCallContract(t, client) },
+		}, {
+			"testCodeOverride",
+			func(t *testing.T) { testCodeOverride(t) },
+		}, {
+			"testNoCodeOverride",
+			func(t *testing.T) { testNoCodeOverride(t) },
 		},
 	}
 	t.Parallel()
@@ -327,3 +335,78 @@ func testCallContract(t *testing.T, client *rpc.Client) {
 		t.Fatalf("unexpected error: %v", err)
 	}
 }
+
+func testNoCodeOverride(t *testing.T) {
+	override := OverrideAccount{}
+
+	mapAcc := make(map[common.Address]OverrideAccount)
+	mapAcc[testAddr] = override
+
+	encoded := toOverrideMap(&mapAcc)
+
+	type overrideAccount struct {
+		Nonce     hexutil.Uint64              `json:"nonce"`
+		Code      []byte                      `json:"code"`
+		Balance   *hexutil.Big                `json:"balance"`
+		State     map[common.Hash]common.Hash `json:"state"`
+		StateDiff map[common.Hash]common.Hash `json:"stateDiff"`
+	}
+
+	var accounts map[common.Address]*overrideAccount
+	marshalled, err := json.Marshal(encoded)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	err = json.Unmarshal(marshalled, &accounts)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	account := accounts[testAddr]
+	if account == nil {
+		t.Fatal("Empty account retrieved")
+	}
+
+	if account.Code != nil {
+		t.Fatal("Code should be empty")
+	}
+}
+
+func testCodeOverride(t *testing.T) {
+	override := OverrideAccount{}
+	override.Code = []byte{1}
+
+	mapAcc := make(map[common.Address]OverrideAccount)
+	mapAcc[testAddr] = override
+
+	encoded := toOverrideMap(&mapAcc)
+
+	type overrideAccount struct {
+		Nonce     hexutil.Uint64              `json:"nonce"`
+		Code      []byte                      `json:"code"`
+		Balance   *hexutil.Big                `json:"balance"`
+		State     map[common.Hash]common.Hash `json:"state"`
+		StateDiff map[common.Hash]common.Hash `json:"stateDiff"`
+	}
+
+	var accounts map[common.Address]*overrideAccount
+	marshalled, err := json.Marshal(encoded)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	err = json.Unmarshal(marshalled, &accounts)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	account := accounts[testAddr]
+	if account == nil {
+		t.Fatal("Empty account retrieved")
+	}
+
+	if account.Code == nil {
+		t.Fatal("Code should not be empty")
+	}
+}