Use shadowing to avoid relisting chainstate fields

cosmos · Sep 21, 2023 · f8c8329 · f8c8329
1 parent 05232d9
commit f8c8329
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 65 deletions.
diff --git a/tests/e2e/json_utils.go b/tests/e2e/json_utils.go
@@ -12,21 +12,16 @@ type ProposalAndType struct {
 	Type        string
 }
 
-// duplicated from the ChainState with a minor change to the Proposals field
-type ChainStateWithProposalTypes struct {
-	ValBalances                    *map[ValidatorID]uint
-	ValPowers                      *map[ValidatorID]uint
-	RepresentativePowers           *map[ValidatorID]uint
-	Params                         *[]Param
-	Rewards                        *Rewards
-	ConsumerChains                 *map[ChainID]bool
-	AssignedKeys                   *map[ValidatorID]string
-	ProviderKeys                   *map[ValidatorID]string
-	ConsumerChainQueueSizes        *map[ChainID]uint
-	GlobalSlashQueueSize           *uint
-	RegisteredConsumerRewardDenoms *[]string
-	Proposals                      *map[uint]ProposalAndType // the only thing changed from the real ChainState
-}
+type (
+	// to have a ChainState object that does not have the overriden Marshal/Unmarshal method
+	ChainStateCopy ChainState
+
+	// duplicated from the ChainState with a minor change to the Proposals field
+	ChainStateWithProposalTypes struct {
+		ChainStateCopy
+		Proposals *map[uint]ProposalAndType // the only thing changed from the real ChainState
+	}
+)
 
 // MarshalJSON marshals a step into JSON while including the type of the action.
 func (step Step) MarshalJSON() ([]byte, error) {
@@ -52,6 +47,7 @@ func (step *Step) UnmarshalJSON(data []byte) error {
 		Action     json.RawMessage
 		State      State
 	}
+	print(string(data))
 	if err := json.Unmarshal(data, &tmp); err != nil {
 		return err
 	}
@@ -290,19 +286,8 @@ func UnmarshalMapToActionType(rawAction json.RawMessage, actionTypeString string
 
 // MarshalJSON transforms the ChainState into a ChainStateWithProposalTypes by adding type info to the proposals
 func (c ChainState) MarshalJSON() ([]byte, error) {
-	chainStateWithProposalTypes := ChainStateWithProposalTypes{
-		ValBalances:                    c.ValBalances,
-		ValPowers:                      c.ValPowers,
-		RepresentativePowers:           c.RepresentativePowers,
-		Params:                         c.Params,
-		Rewards:                        c.Rewards,
-		ConsumerChains:                 c.ConsumerChains,
-		AssignedKeys:                   c.AssignedKeys,
-		ProviderKeys:                   c.ProviderKeys,
-		ConsumerChainQueueSizes:        c.ConsumerChainQueueSizes,
-		GlobalSlashQueueSize:           c.GlobalSlashQueueSize,
-		RegisteredConsumerRewardDenoms: c.RegisteredConsumerRewardDenoms,
-	}
+	chainStateCopy := ChainStateCopy(c)
+	chainStateWithProposalTypes := ChainStateWithProposalTypes{chainStateCopy, nil}
 	if c.Proposals != nil {
 		proposalsWithTypes := make(map[uint]ProposalAndType)
 		for k, v := range *c.Proposals {
@@ -324,17 +309,9 @@ func (c *ChainState) UnmarshalJSON(data []byte) error {
 	if err != nil {
 		return err
 	}
-	c.ValBalances = chainStateWithProposalTypes.ValBalances
-	c.ValPowers = chainStateWithProposalTypes.ValPowers
-	c.RepresentativePowers = chainStateWithProposalTypes.RepresentativePowers
-	c.Params = chainStateWithProposalTypes.Params
-	c.Rewards = chainStateWithProposalTypes.Rewards
-	c.ConsumerChains = chainStateWithProposalTypes.ConsumerChains
-	c.AssignedKeys = chainStateWithProposalTypes.AssignedKeys
-	c.ProviderKeys = chainStateWithProposalTypes.ProviderKeys
-	c.ConsumerChainQueueSizes = chainStateWithProposalTypes.ConsumerChainQueueSizes
-	c.GlobalSlashQueueSize = chainStateWithProposalTypes.GlobalSlashQueueSize
-	c.RegisteredConsumerRewardDenoms = chainStateWithProposalTypes.RegisteredConsumerRewardDenoms
+
+	chainState := ChainState(chainStateWithProposalTypes.ChainStateCopy)
+	*c = chainState
 
 	if chainStateWithProposalTypes.Proposals != nil {
 		proposals := make(map[uint]Proposal)

diff --git a/tests/e2e/trace_handlers_test.go b/tests/e2e/trace_handlers_test.go
@@ -62,33 +62,14 @@ func TestWriterThenParser(t *testing.T) {
 	for name, tc := range tests {
 		t.Run(name, func(t *testing.T) {
 			filename := filepath.Join(dir, "trace.json")
-			err := WriteReadCompareTrace(tc.trace, filename, name)
+			err := WriteAndReadTrace(GlobalJSONParser, GlobalJSONWriter, tc.trace, filename)
 			if err != nil {
-				log.Fatal(err)
+				log.Fatalf("got error for testcase %v: %s", name, err)
 			}
 		})
 	}
 }
 
-// Write a trace to a file, then reads it back and compares to the original.
-func WriteReadCompareTrace(trace []Step, filename, name string) error {
-	err := WriteAndReadTrace(GlobalJSONParser, GlobalJSONWriter, trace, filename)
-	if err != nil {
-		return fmt.Errorf("in testcase %v, got error writing trace to file: %v", name, err)
-	}
-
-	got, err := GlobalJSONParser.ReadTraceFromFile(filename)
-	if err != nil {
-		return fmt.Errorf("in testcase %v, got error reading trace from file: %v", name, err)
-	}
-	diff := cmp.Diff(trace, got, cmp.AllowUnexported(Step{}))
-	if diff != "" {
-		return fmt.Errorf("Got a difference for testcase %s (-want +got):\n%s", name, diff)
-	}
-
-	return nil
-}
-
 // Checks that writing a trace does not result in an error.
 func TestWriteExamples(t *testing.T) {
 	tests := map[string]struct {
@@ -236,15 +217,13 @@ func WriteAndReadTrace(parser TraceParser, writer TraceWriter, trace []Step, tmp
 		return fmt.Errorf("error writing trace to file: %v", err)
 	}
 
-	got, err := parser.ReadTraceFromFile(tmp_filepath)
+	got, err := GlobalJSONParser.ReadTraceFromFile(tmp_filepath)
 	if err != nil {
-		return fmt.Errorf("error reading trace from file: %v", err)
+		return fmt.Errorf("got error reading trace from file: %v", err)
 	}
-
 	diff := cmp.Diff(trace, got, cmp.AllowUnexported(Step{}))
 	if diff != "" {
-		return fmt.Errorf(diff)
+		return fmt.Errorf("Got a difference (-want +got):\n%s", diff)
 	}
-
 	return nil
 }