From aeeb4cbcd069507c101038101d12ee1c7373fd52 Mon Sep 17 00:00:00 2001
From: wonjoon <wnjoon@gmail.com>
Date: Tue, 14 Jan 2025 02:47:19 +0900
Subject: [PATCH] chore: test

---
 test/e2e/configurer/upgrade.go    | 14 ++++++++++++++
 test/e2e/initialization/export.go |  3 +++
 test/e2e/initialization/node.go   | 18 ++++++++++++------
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/test/e2e/configurer/upgrade.go b/test/e2e/configurer/upgrade.go
index 123a3924f..c2c45b418 100644
--- a/test/e2e/configurer/upgrade.go
+++ b/test/e2e/configurer/upgrade.go
@@ -22,6 +22,9 @@ import (
 	"github.com/babylonlabs-io/babylon/test/e2e/configurer/config"
 	"github.com/babylonlabs-io/babylon/test/e2e/containers"
 	"github.com/babylonlabs-io/babylon/test/e2e/initialization"
+	"github.com/babylonlabs-io/babylon/testutil/signer"
+	cmtCfg "github.com/cometbft/cometbft/config"
+	cmtos "github.com/cometbft/cometbft/libs/os"
 )
 
 type UpgradeSettings struct {
@@ -257,6 +260,9 @@ func (uc *UpgradeConfigurer) upgradeContainers(chainConfig *chain.Config, propHe
 	uc.containerManager.CurrentTag = "latest"
 
 	for _, node := range chainConfig.NodeConfigs {
+		if err := saveIfNotExists(node); err != nil {
+			return err
+		}
 		if err := node.Run(); err != nil {
 			return err
 		}
@@ -268,6 +274,14 @@ func (uc *UpgradeConfigurer) upgradeContainers(chainConfig *chain.Config, propHe
 	return nil
 }
 
+func saveIfNotExists(cfg *chain.NodeConfig) error {
+	nodeDir := cfg.Node.ConfigDir
+	if !cmtos.FileExists(filepath.Join(nodeDir, cmtCfg.DefaultConfig().PrivValidatorKeyFile())) {
+		return signer.GeneratePrivSigner(nodeDir)
+	}
+	return nil
+}
+
 // ParseGovPropFromFile loads the proposal from the UpgradeSignetLaunchFilePath
 func (uc *UpgradeConfigurer) ParseGovPropFromFile() (*upgradetypes.MsgSoftwareUpgrade, error) {
 	pwd, err := os.Getwd()
diff --git a/test/e2e/initialization/export.go b/test/e2e/initialization/export.go
index 6c741a4c0..e4ab0622a 100644
--- a/test/e2e/initialization/export.go
+++ b/test/e2e/initialization/export.go
@@ -2,6 +2,8 @@ package initialization
 
 import (
 	"fmt"
+
+	"github.com/babylonlabs-io/babylon/privval"
 )
 
 type ChainMeta struct {
@@ -19,6 +21,7 @@ type Node struct {
 	PrivateKey    []byte `json:"privateKey"`
 	PeerId        string `json:"peerId"`
 	IsValidator   bool   `json:"isValidator"`
+	ConsensusKey  privval.WrappedFilePVKey
 }
 
 type Chain struct {
diff --git a/test/e2e/initialization/node.go b/test/e2e/initialization/node.go
index d148c0d00..0dfae6766 100644
--- a/test/e2e/initialization/node.go
+++ b/test/e2e/initialization/node.go
@@ -34,6 +34,7 @@ import (
 	babylonApp "github.com/babylonlabs-io/babylon/app"
 	appparams "github.com/babylonlabs-io/babylon/app/params"
 	"github.com/babylonlabs-io/babylon/cmd/babylond/cmd"
+	"github.com/babylonlabs-io/babylon/crypto/bls12381"
 	"github.com/babylonlabs-io/babylon/privval"
 	"github.com/babylonlabs-io/babylon/test/e2e/util"
 	cmtprivval "github.com/cometbft/cometbft/privval"
@@ -180,20 +181,24 @@ func (n *internalNode) createConsensusKey() error {
 
 	// create new key for consensus
 	// file pv
-	var privKey ed25519.PrivKey
+	var cmtPrivKey ed25519.PrivKey
+	var blsPrivKey bls12381.PrivateKey
 	if n.mnemonic == "" {
-		privKey = ed25519.GenPrivKey()
+		cmtPrivKey = ed25519.GenPrivKey()
+		blsPrivKey = bls12381.GenPrivKey()
 	} else {
-		privKey = ed25519.GenPrivKeyFromSecret([]byte(n.mnemonic))
+		cmtPrivKey = ed25519.GenPrivKeyFromSecret([]byte(n.mnemonic))
+		blsPrivKey = bls12381.GenPrivKeyFromSecret([]byte(n.mnemonic))
 	}
-	filePV := cmtprivval.NewFilePV(privKey, pvKeyFile, pvStateFile)
+
+	filePV := cmtprivval.NewFilePV(cmtPrivKey, pvKeyFile, pvStateFile)
 	filePV.Key.Save()
 	filePV.LastSignState.Save()
 
 	// bls pv
-	blsPV := privval.GenBlsPV(blsKeyFile, blsPasswordFile, "password", accAddress.String())
+	blsPV := privval.NewBlsPV(blsPrivKey, blsKeyFile, blsPasswordFile, accAddress.String())
+	blsPV.Key.Save("password", accAddress.String())
 
-	// n.consensusKey = filePV.Key
 	n.consensusKey = privval.WrappedFilePVKey{
 		CometPVKey: filePV.Key,
 		BlsPVKey:   blsPV.Key,
@@ -267,6 +272,7 @@ func (n *internalNode) export() *Node {
 		PrivateKey:    n.privateKey.Bytes(),
 		PeerId:        n.peerId,
 		IsValidator:   n.isValidator,
+		ConsensusKey:  n.consensusKey,
 	}
 }