- update comments

- add e2e test for wallet creation result replay

Author: johnnynanjiang

e2e/keygen_test.go

@@ -79,10 +79,23 @@ func testKeyGeneration(t *testing.T, suite *E2ETestSuite) {
 	if suite.mpcClient == nil {
 		t.Fatal("MPC client is not initialized. Make sure Setup subtest runs first.")
 	}
+
+	walletID := uuid.New().String()
+
+	// Create a new wallet
+	createWalletAndVerify(t, suite, walletID)
+
+	// Create another new wallet with the same wallet ID
+	createWalletAndVerify(t, suite, walletID)
+
+	t.Log("Key generation test completed")
+}
+
+func createWalletAndVerify(t *testing.T, suite *E2ETestSuite, walletID string) {
 	// Generate 1 wallet ID for testing
 	walletIDs := make([]string, 0, 10)
 	for i := 0; i < 1; i++ {
-		walletIDs = append(walletIDs, uuid.New().String())
+		walletIDs = append(walletIDs, walletID)
 		suite.walletIDs = append(suite.walletIDs, walletIDs[i])
 	}
 
@@ -92,7 +105,11 @@ func testKeyGeneration(t *testing.T, suite *E2ETestSuite) {
 	err := suite.mpcClient.OnWalletCreationResult(func(result event.KeygenResultEvent) {
 		logger.Info("On wallet creation result", "event", result)
 		t.Logf("Received keygen result for wallet %s: %s", result.WalletID, result.ResultType)
-		suite.keygenResults[result.WalletID] = &result
+
+		// For testing replay properly, don't overwrite existing result for the same wallet ID to preserve the first result
+		if _, exists := suite.keygenResults[result.WalletID]; !exists {
+			suite.keygenResults[result.WalletID] = &result
+		}
 
 		if result.ResultType == event.ResultTypeError {
 			t.Logf("Keygen failed for wallet %s: %s (%s)", result.WalletID, result.ErrorReason, result.ErrorCode)
@@ -165,8 +182,6 @@ checkResults:
 			assert.NotEmpty(t, result.EDDSAPubKey, "EdDSA public key should not be empty for wallet %s", walletID)
 		}
 	}
-
-	t.Log("Key generation test completed")
 }
 
 func verifyKeyConsistency(t *testing.T, suite *E2ETestSuite) {

pkg/eventconsumer/event_consumer.go

@@ -159,6 +159,7 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 
 	walletID := msg.WalletID
 
+	// Attempt to get previously stored wallet creation result (if any) by the wallet ID
 	storedWalletCreationResult, storedWalletCreationResultError := ec.node.GetWalletCreationResult(walletID)
 
 	// Error when retrieving wallet creation result for the wallet ID
@@ -168,12 +169,8 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 	}
 
 	// Replay the wallet creation result if it already exists for the wallet ID
-	// TODO
-	// 1. to inspect storedWalletCreationResult
-	// 2. to move the following duplicate logic (line 341~351) into a func
+	// TODO: to move the following duplicate logic (line 308 and line 341) into a func
 	if storedWalletCreationResult != nil {
-		logger.Info("storedWalletCreationResult", storedWalletCreationResult)
-
 		key := event.KeygenResultSubject(natMsg.Header.Get(event.ClientIDHeader), walletID)
 		if err := ec.genKeyResultQueue.Enqueue(key, storedWalletCreationResult, &messaging.EnqueueOptions{
 			IdempotententKey: composeKeygenIdempotentKey(walletID, natMsg),
@@ -302,7 +299,6 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 	}
 
 	// Store wallet creation result
-	logger.Info("payload", payload)
 	if storeErr := ec.node.StoreWalletCreationResult(walletID, payload); storeErr != nil {
 		logger.Error("Failed to store wallet creation result", storeErr, "walletID", walletID)
 		ec.handleKeygenSessionError(walletID, storeErr, "Failed to store wallet creation result", natMsg)