fix: e2e trace format fails on slashthrottlesteps (#903)

* Add other steps to read/write test

* Improve error logging in trace handler test

* Push ActionType handling into Unmarshal/Marshal and out of own class

* Rename generic trace handler files to be clear they are for json

* Rename to marshalling to be more generic

* Add marshal/unmarshal for proposals

* Export unexported field

* Add test for marshal/unmarshalling chain state

* Fix pointer issues

* Fix typo: action -> proposal

* Log proposal string in test

* Use json.RawMessage instead of map[string]any

* For uniformity, also use RawMessage for step unmarshalling

* Add tests for extra proposal types

* Add more proposal types to test and unify names

* Add handling for ParamsProposal

* Regenerate traces

* Chore: Export forgotten field

* Use string, not int, to help marshal/unmarshal

tests/e2e/json_chainState_marshalling.go

@@ -0,0 +1,168 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"reflect"
+)
+
+type ProposalWithType struct {
+	ProposalType string
+	Proposal     Proposal
+}
+
+// MarshalJSON marshals a chainState into JSON while including the type of the proposal.
+func (chainState ChainState) MarshalJSON() ([]byte, error) {
+	var proposalsWithTypes map[uint]ProposalWithType
+	if chainState.Proposals != nil {
+		proposalsWithTypes = make(map[uint]ProposalWithType, len(*chainState.Proposals))
+
+		for k, v := range *chainState.Proposals {
+			proposalsWithTypes[k] = ProposalWithType{
+				ProposalType: reflect.TypeOf(v).String(),
+				Proposal:     v,
+			}
+		}
+	} else {
+		proposalsWithTypes = nil
+	}
+
+	result := struct {
+		ValBalances             *map[ValidatorID]uint
+		Proposals               *map[uint]ProposalWithType
+		ValPowers               *map[ValidatorID]uint
+		RepresentativePowers    *map[ValidatorID]uint
+		Params                  *[]Param
+		Rewards                 *Rewards
+		ConsumerChains          *map[ChainID]bool
+		AssignedKeys            *map[ValidatorID]string
+		ProviderKeys            *map[ValidatorID]string // validatorID: validator provider key
+		ConsumerChainQueueSizes *map[ChainID]uint
+		GlobalSlashQueueSize    *uint
+	}{
+		ValBalances:             chainState.ValBalances,
+		Proposals:               &proposalsWithTypes,
+		ValPowers:               chainState.ValPowers,
+		RepresentativePowers:    chainState.RepresentativePowers,
+		Params:                  chainState.Params,
+		Rewards:                 chainState.Rewards,
+		ConsumerChains:          chainState.ConsumerChains,
+		AssignedKeys:            chainState.AssignedKeys,
+		ProviderKeys:            chainState.ProviderKeys,
+		ConsumerChainQueueSizes: chainState.ConsumerChainQueueSizes,
+		GlobalSlashQueueSize:    chainState.GlobalSlashQueueSize,
+	}
+
+	return json.Marshal(result)
+}
+
+func (state *ChainState) UnmarshalJSON(data []byte) error {
+	var tmp struct {
+		ValBalances             *map[ValidatorID]uint
+		Proposals               *map[uint]json.RawMessage
+		ValPowers               *map[ValidatorID]uint
+		RepresentativePowers    *map[ValidatorID]uint
+		Params                  *[]Param
+		Rewards                 *Rewards
+		ConsumerChains          *map[ChainID]bool
+		AssignedKeys            *map[ValidatorID]string
+		ProviderKeys            *map[ValidatorID]string // validatorID: validator provider key
+		ConsumerChainQueueSizes *map[ChainID]uint
+		GlobalSlashQueueSize    *uint
+	}
+
+	err := json.Unmarshal(data, &tmp)
+	if err != nil {
+		return err
+	}
+
+	var proposals *map[uint]Proposal
+	if tmp.Proposals != nil {
+		proposals, err = UnmarshalProposals(*tmp.Proposals)
+		if err != nil {
+			return err
+		}
+	}
+
+	state.Proposals = proposals
+
+	state.ValBalances = tmp.ValBalances
+	state.ValPowers = tmp.ValPowers
+	state.RepresentativePowers = tmp.RepresentativePowers
+	state.Params = tmp.Params
+	state.Rewards = tmp.Rewards
+	state.ConsumerChains = tmp.ConsumerChains
+	state.AssignedKeys = tmp.AssignedKeys
+	state.ProviderKeys = tmp.ProviderKeys
+	state.ConsumerChainQueueSizes = tmp.ConsumerChainQueueSizes
+	state.GlobalSlashQueueSize = tmp.GlobalSlashQueueSize
+
+	return nil
+}
+
+func UnmarshalProposals(proposals map[uint]json.RawMessage) (*map[uint]Proposal, error) {
+	result := make(map[uint]Proposal, len(proposals))
+
+	for k, v := range proposals {
+		var tmp struct {
+			Proposal     json.RawMessage
+			ProposalType string
+		}
+
+		if err := json.Unmarshal(v, &tmp); err != nil {
+			return nil, err
+		}
+
+		proposal, err := UnmarshalMapToProposalType(tmp.Proposal, tmp.ProposalType)
+		if err != nil {
+			return nil, err
+		}
+		result[k] = proposal
+	}
+
+	return &result, nil
+}
+
+// UnmarshalMapToProposalType takes a JSON message and a proposal type and marshals into an object of the corresponding proposal.
+func UnmarshalMapToProposalType(input json.RawMessage, proposalType string) (Proposal, error) {
+	switch proposalType {
+	case "main.ConsumerAdditionProposal":
+		var proposal ConsumerAdditionProposal
+		err := json.Unmarshal(input, &proposal)
+		if err != nil {
+			return nil, err
+		}
+		return proposal, nil
+	case "main.ConsumerRemovalProposal":
+		var proposal ConsumerRemovalProposal
+		err := json.Unmarshal(input, &proposal)
+		if err != nil {
+			return nil, err
+		}
+		return proposal, nil
+	case "main.EquivocationProposal":
+		var proposal EquivocationProposal
+		err := json.Unmarshal(input, &proposal)
+		if err != nil {
+			return nil, err
+		}
+		return proposal, nil
+	case "main.ParamsProposal":
+		var proposal ParamsProposal
+		err := json.Unmarshal(input, &proposal)
+		if err != nil {
+			return nil, err
+		}
+		return proposal, nil
+
+	case "main.TextProposal":
+		var proposal TextProposal
+		err := json.Unmarshal(input, &proposal)
+		if err != nil {
+			return nil, err
+		}
+		return proposal, nil
+	default:
+		return nil, fmt.Errorf("%s is not a known proposal type", proposalType)
+	}
+}

tests/e2e/trace_parser.go → tests/e2e/json_parser.go

@@ -3,7 +3,6 @@ package main
 import (
 	"bytes"
 	"encoding/json"
-	"fmt"
 	"os"
 	"path/filepath"
 )
@@ -24,29 +23,19 @@ func (parser JSONParser) ReadTraceFromFile(path string) ([]Step, error) {
 	}
 
 	// Unmarshal the JSON into a slice of Step structs
+<<<<<<< HEAD:tests/e2e/trace_parser.go
 	var stepsWithActionTypes []StepWithActionType
 
 	decoder := json.NewDecoder(bytes.NewReader(jsonData))
 	decoder.DisallowUnknownFields() // To avoid silent errors. Will cause an error if the JSON contains unknown fields
 	err = decoder.Decode(&stepsWithActionTypes)
+=======
+	var steps []Step
+	err = json.Unmarshal(jsonData, &steps)
+>>>>>>> 24b0ef1b (fix: e2e trace format fails on slashthrottlesteps (#903)):tests/e2e/json_parser.go
 	if err != nil {
 		return nil, err
 	}
 
-	steps := make([]Step, len(stepsWithActionTypes))
-
-	// Unmarshal the actions inside the steps from map[string]any to the corresponding action type
-	for i, step := range stepsWithActionTypes {
-		action, err := UnmarshalMapToActionType(step.Action.(map[string]any), step.ActionType)
-		if err != nil {
-			return nil, err
-		}
-
-		fmt.Println(action)
-
-		steps[i] = Step{action, step.State}
-
-	}
-
 	return steps, nil
 }

tests/e2e/trace_utils.go → tests/e2e/json_step_marshalling.go

@@ -4,14 +4,49 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+<<<<<<< HEAD:tests/e2e/trace_utils.go
 
 	"github.com/mitchellh/mapstructure"
+=======
+>>>>>>> 24b0ef1b (fix: e2e trace format fails on slashthrottlesteps (#903)):tests/e2e/json_step_marshalling.go
 )
 
-// StepWithActionType is a utility class that wraps a Step together with its action type. Used to marshal/unmarshal
-type StepWithActionType struct {
-	Step
-	ActionType string `json:"ActionType"`
+// MarshalJSON marshals a step into JSON while including the type of the action.
+func (step Step) MarshalJSON() ([]byte, error) {
+	actionType := reflect.TypeOf(step.Action).String()
+
+	result := struct {
+		ActionType string
+		Action     interface{}
+		State      State
+	}{
+		ActionType: actionType,
+		Action:     step.Action,
+		State:      step.State,
+	}
+
+	return json.Marshal(result)
+}
+
+// UnmarshalJSON unmarshals a step from JSON while including the type of the action.
+func (step *Step) UnmarshalJSON(data []byte) error {
+	var tmp struct {
+		ActionType string
+		Action     json.RawMessage
+		State      State
+	}
+	if err := json.Unmarshal(data, &tmp); err != nil {
+		return err
+	}
+
+	action, err := UnmarshalMapToActionType(tmp.Action, tmp.ActionType)
+	if err != nil {
+		return err
+	}
+
+	step.Action = action
+	step.State = tmp.State
+	return nil
 }
 
 // has to be manually kept in sync with the available action types.

tests/e2e/steps_democracy.go

@@ -67,9 +67,15 @@ func stepsDemocracy(consumerName string) []Step {
 				Chain:    ChainID(consumerName),
 				From:     ValidatorID("alice"),
 				Deposit:  10000001,
+<<<<<<< HEAD
 				Subspace: "transfer",
 				Key:      "SendEnabled",
 				Value:    true,
+=======
+				Subspace: "staking",
+				Key:      "MaxValidators",
+				Value:    "105",
+>>>>>>> 24b0ef1b (fix: e2e trace format fails on slashthrottlesteps (#903))
 			},
 			State: State{
 				ChainID(consumerName): ChainState{