From b845cbb91f78eaaa7f27ad52ab5015f7b88917e2 Mon Sep 17 00:00:00 2001
From: Yaru Wang <yaru@informal.systems>
Date: Mon, 26 Jun 2023 10:52:25 +0200
Subject: [PATCH] test: fix test TestInitAndExportGenesis

---
 x/ccv/provider/keeper/genesis_test.go | 58 +++++++++++++++++----------
 1 file changed, 36 insertions(+), 22 deletions(-)

diff --git a/x/ccv/provider/keeper/genesis_test.go b/x/ccv/provider/keeper/genesis_test.go
index 1b56c4a79f..bf327680d1 100644
--- a/x/ccv/provider/keeper/genesis_test.go
+++ b/x/ccv/provider/keeper/genesis_test.go
@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"sort"
 	"testing"
 	"time"
 
@@ -35,8 +36,22 @@ func TestInitAndExportGenesis(t *testing.T) {
 	consumerTmPubKey := consumerCryptoId.TMProtoCryptoPublicKey()
 	consumerConsAddr := consumerCryptoId.ConsumerConsAddress()
 
-	initTimeoutTime := uint64(time.Now().UTC().UnixNano()) + 10
-	vscSendTime := time.Now().Add(-30 * time.Minute)
+	initTimeoutTimeStamps := []providertypes.InitTimeoutTimestamp{
+		{ChainId: cChainIDs[0], Timestamp: uint64(time.Now().UTC().UnixNano()) + 10},
+		{ChainId: cChainIDs[1], Timestamp: uint64(time.Now().UTC().UnixNano()) + 15},
+	}
+
+	now := time.Now().UTC()
+	vscSendTimeStampsC0 := []providertypes.VscSendTimestamp{
+		{ChainId: "c0", VscId: 1, Timestamp: now.Add(time.Hour)},
+		{ChainId: "c0", VscId: 2, Timestamp: now.Add(2 * time.Hour)},
+	}
+
+	vscSendTimeStampsC1 := []providertypes.VscSendTimestamp{
+		{ChainId: "c1", VscId: 1, Timestamp: now.Add(-time.Hour)},
+		{ChainId: "c1", VscId: 2, Timestamp: now.Add(time.Hour)},
+	}
+	vscSendTimeStampsAll := append(vscSendTimeStampsC0, vscSendTimeStampsC1...)
 	// create genesis struct
 	provGenesis := providertypes.NewGenesisState(vscID,
 		[]providertypes.ValsetUpdateIdToHeight{{ValsetUpdateId: vscID, Height: initHeight}},
@@ -99,19 +114,8 @@ func TestInitAndExportGenesis(t *testing.T) {
 				ConsumerAddrs: &providertypes.AddressList{Addresses: [][]byte{consumerConsAddr.ToSdkConsAddr()}},
 			},
 		},
-		[]providertypes.InitTimeoutTimestamp{
-			{
-				ChainId:   cChainIDs[0],
-				Timestamp: initTimeoutTime,
-			},
-		},
-		[]providertypes.VscSendTimestamp{
-			{
-				ChainId:   cChainIDs[0],
-				VscId:     vscID,
-				Timestamp: vscSendTime,
-			},
-		},
+		initTimeoutTimeStamps,
+		vscSendTimeStampsAll,
 	)
 
 	// Instantiate in-mem provider keeper with mocks
@@ -179,13 +183,23 @@ func TestInitAndExportGenesis(t *testing.T) {
 	// check the exported genesis
 	require.Equal(t, provGenesis, pk.ExportGenesis(ctx))
 
-	initTimeoutTimestamp, found := pk.GetInitTimeoutTimestamp(ctx, cChainIDs[0])
-	require.Equal(t, initTimeoutTimestamp, initTimeoutTime)
-	require.True(t, found)
-
-	vscSendTimestamp, found := pk.GetVscSendTimestamp(ctx, cChainIDs[0], vscID)
-	require.Equal(t, vscSendTimestamp, vscSendTime)
-	require.True(t, found)
+	initTimeoutTimestampInStore := pk.GetAllInitTimeoutTimestamps(ctx)
+	sort.Slice(initTimeoutTimestampInStore, func(i, j int) bool {
+		return initTimeoutTimestampInStore[i].Timestamp < initTimeoutTimestampInStore[j].Timestamp
+	})
+	require.Equal(t, initTimeoutTimestampInStore, initTimeoutTimeStamps)
+
+	vscSendTimestampsC0InStore := pk.GetAllVscSendTimestamps(ctx, cChainIDs[0])
+	sort.Slice(vscSendTimestampsC0InStore, func(i, j int) bool {
+		return vscSendTimestampsC0InStore[i].VscId < vscSendTimestampsC0InStore[j].VscId
+	})
+	require.Equal(t, vscSendTimestampsC0InStore, vscSendTimeStampsC0)
+
+	vscSendTimestampsC1InStore := pk.GetAllVscSendTimestamps(ctx, cChainIDs[1])
+	sort.Slice(vscSendTimestampsC1InStore, func(i, j int) bool {
+		return vscSendTimestampsC1InStore[i].VscId < vscSendTimestampsC1InStore[j].VscId
+	})
+	require.Equal(t, vscSendTimestampsC1InStore, vscSendTimeStampsC1)
 }
 
 func assertConsumerChainStates(t *testing.T, ctx sdk.Context, pk keeper.Keeper, consumerStates ...providertypes.ConsumerState) {