From 2743e6492c920f5bb48945fd3f65464bf3540ac3 Mon Sep 17 00:00:00 2001
From: Tomasz Slabon <tomasz.slabon@keep.network>
Date: Tue, 11 Jun 2024 15:32:27 +0200
Subject: [PATCH] Set wallet data in unit tests

---
 pkg/tbtc/chain_test.go      | 21 ++++++++++++++----
 pkg/tbtc/inactivity_test.go | 10 ++++++---
 pkg/tbtc/node_test.go       | 43 +++++++++++++++++++++++++++++++++++++
 pkg/tbtc/signing_test.go    | 15 +++++++++++++
 4 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/pkg/tbtc/chain_test.go b/pkg/tbtc/chain_test.go
index d1ce5775c3..15bb4c94ca 100644
--- a/pkg/tbtc/chain_test.go
+++ b/pkg/tbtc/chain_test.go
@@ -843,10 +843,6 @@ func buildDepositRequestKey(
 	return sha256.Sum256(append(fundingTxHash[:], buffer...))
 }
 
-func (lc *localChain) IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error) {
-	panic("unsupported")
-}
-
 func (lc *localChain) CalculateWalletID(
 	walletPublicKey *ecdsa.PublicKey,
 ) ([32]byte, error) {
@@ -896,6 +892,23 @@ func (lc *localChain) GetWallet(walletPublicKeyHash [20]byte) (
 	return walletChainData, nil
 }
 
+func (lc *localChain) IsWalletRegistered(EcdsaWalletID [32]byte) (bool, error) {
+	lc.walletsMutex.Lock()
+	defer lc.walletsMutex.Unlock()
+
+	for _, walletData := range lc.wallets {
+		if EcdsaWalletID == walletData.EcdsaWalletID {
+			if walletData.State == StateClosed ||
+				walletData.State == StateTerminated {
+				return false, nil
+			}
+			return true, nil
+		}
+	}
+
+	return false, fmt.Errorf("wallet not found")
+}
+
 func (lc *localChain) setWallet(
 	walletPublicKeyHash [20]byte,
 	walletChainData *WalletChainData,
diff --git a/pkg/tbtc/inactivity_test.go b/pkg/tbtc/inactivity_test.go
index b1065e4607..3c8370e0d5 100644
--- a/pkg/tbtc/inactivity_test.go
+++ b/pkg/tbtc/inactivity_test.go
@@ -157,12 +157,16 @@ func setupInactivityClaimExecutorScenario(t *testing.T) (
 	keyStorePersistence := createMockKeyStorePersistence(t, signers...)
 
 	walletPublicKeyHash := bitcoin.PublicKeyHash(signers[0].wallet.publicKey)
-	ecdsaWalletID := [32]byte{1, 2, 3}
+	walletID, err := localChain.CalculateWalletID(signers[0].wallet.publicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	localChain.setWallet(
 		walletPublicKeyHash,
 		&WalletChainData{
-			EcdsaWalletID: ecdsaWalletID,
+			EcdsaWalletID: walletID,
+			State:         StateLive,
 		},
 	)
 
@@ -191,7 +195,7 @@ func setupInactivityClaimExecutorScenario(t *testing.T) (
 		t.Fatal("node is supposed to control wallet signers")
 	}
 
-	return executor, ecdsaWalletID, localChain
+	return executor, walletID, localChain
 }
 
 func TestSignClaim_SigningSuccessful(t *testing.T) {
diff --git a/pkg/tbtc/node_test.go b/pkg/tbtc/node_test.go
index b9dbb01992..bedfb30995 100644
--- a/pkg/tbtc/node_test.go
+++ b/pkg/tbtc/node_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/keep-network/keep-common/pkg/persistence"
 	"github.com/keep-network/keep-core/internal/testutils"
+	"github.com/keep-network/keep-core/pkg/bitcoin"
 	"github.com/keep-network/keep-core/pkg/chain"
 	"github.com/keep-network/keep-core/pkg/generator"
 	"github.com/keep-network/keep-core/pkg/internal/tecdsatest"
@@ -31,6 +32,20 @@ func TestNode_GetSigningExecutor(t *testing.T) {
 
 	signer := createMockSigner(t)
 
+	walletPublicKeyHash := bitcoin.PublicKeyHash(signer.wallet.publicKey)
+	walletID, err := localChain.CalculateWalletID(signer.wallet.publicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	localChain.setWallet(
+		walletPublicKeyHash,
+		&WalletChainData{
+			EcdsaWalletID: walletID,
+			State:         StateLive,
+		},
+	)
+
 	// Populate the mock keystore with the mock signer's data. This is
 	// required to make the node controlling the signer's wallet.
 	keyStorePersistence := createMockKeyStorePersistence(t, signer)
@@ -149,6 +164,20 @@ func TestNode_GetCoordinationExecutor(t *testing.T) {
 
 	signer := createMockSigner(t)
 
+	walletPublicKeyHash := bitcoin.PublicKeyHash(signer.wallet.publicKey)
+	walletID, err := localChain.CalculateWalletID(signer.wallet.publicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	localChain.setWallet(
+		walletPublicKeyHash,
+		&WalletChainData{
+			EcdsaWalletID: walletID,
+			State:         StateLive,
+		},
+	)
+
 	// Populate the mock keystore with the mock signer's data. This is
 	// required to make the node controlling the signer's wallet.
 	keyStorePersistence := createMockKeyStorePersistence(t, signer)
@@ -272,6 +301,20 @@ func TestNode_RunCoordinationLayer(t *testing.T) {
 
 	signer := createMockSigner(t)
 
+	walletPublicKeyHash := bitcoin.PublicKeyHash(signer.wallet.publicKey)
+	walletID, err := localChain.CalculateWalletID(signer.wallet.publicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	localChain.setWallet(
+		walletPublicKeyHash,
+		&WalletChainData{
+			EcdsaWalletID: walletID,
+			State:         StateLive,
+		},
+	)
+
 	// Populate the mock keystore with the mock signer's data. This is
 	// required to make the node controlling the signer's wallet.
 	keyStorePersistence := createMockKeyStorePersistence(t, signer)
diff --git a/pkg/tbtc/signing_test.go b/pkg/tbtc/signing_test.go
index 9bcd42ef6c..9298ad7d7f 100644
--- a/pkg/tbtc/signing_test.go
+++ b/pkg/tbtc/signing_test.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/keep-network/keep-core/internal/testutils"
+	"github.com/keep-network/keep-core/pkg/bitcoin"
 	"github.com/keep-network/keep-core/pkg/chain"
 	"github.com/keep-network/keep-core/pkg/chain/local_v1"
 	"github.com/keep-network/keep-core/pkg/generator"
@@ -159,6 +160,20 @@ func setupSigningExecutor(t *testing.T) *signingExecutor {
 		}
 	}
 
+	walletPublicKeyHash := bitcoin.PublicKeyHash(signers[0].wallet.publicKey)
+	walletID, err := localChain.CalculateWalletID(signers[0].wallet.publicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	localChain.setWallet(
+		walletPublicKeyHash,
+		&WalletChainData{
+			EcdsaWalletID: walletID,
+			State:         StateLive,
+		},
+	)
+
 	keyStorePersistence := createMockKeyStorePersistence(t, signers...)
 
 	node, err := newNode(