From ea66925b662f2231631e728cb7ecb61e3eb3f3c9 Mon Sep 17 00:00:00 2001
From: Louis <540788769@qq.com>
Date: Tue, 15 Oct 2024 16:16:42 +0800
Subject: [PATCH] feat: add Optimize the code: Add test case (#2)

---
 wallet/ethereum/address.go      |  2 --
 wallet/ethereum/address_test.go | 58 +++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/wallet/ethereum/address.go b/wallet/ethereum/address.go
index 96131ec..68c071f 100644
--- a/wallet/ethereum/address.go
+++ b/wallet/ethereum/address.go
@@ -3,9 +3,7 @@ package ethereum
 import (
 	"crypto/ecdsa"
 	"encoding/hex"
-
 	"github.com/ethereum/go-ethereum/common"
-
 	"github.com/ethereum/go-ethereum/crypto"
 )
 
diff --git a/wallet/ethereum/address_test.go b/wallet/ethereum/address_test.go
index 137a66d..bec1d6b 100644
--- a/wallet/ethereum/address_test.go
+++ b/wallet/ethereum/address_test.go
@@ -1,7 +1,11 @@
 package ethereum
 
 import (
+	"crypto/elliptic"
+	"encoding/hex"
 	"fmt"
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/crypto"
 	"testing"
 )
 
@@ -12,3 +16,57 @@ func TestCreateAddressByKeyPairs(t *testing.T) {
 	}
 	fmt.Println(address.PrivateKey, address.PublicKey, address.Address)
 }
+
+func TestCreateAddressFromPrivateKey(t *testing.T) {
+	privKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("Failed to generate private key: %v", err)
+	}
+
+	priKeyHex, address, err := CreateAddressFromPrivateKey(privKey)
+	if err != nil {
+		t.Fatalf("Error creating address from private key: %v", err)
+	}
+
+	expectedPriKeyHex := hex.EncodeToString(privKey.D.Bytes())
+	if priKeyHex != expectedPriKeyHex {
+		t.Errorf("Expected private key hex: %s, got: %s", expectedPriKeyHex, priKeyHex)
+	}
+
+	expectedAddress := crypto.PubkeyToAddress(privKey.PublicKey).String()
+	if address != expectedAddress {
+		t.Errorf("Expected address: %s, got: %s", expectedAddress, address)
+	}
+}
+
+func TestPublicKeyToAddress(t *testing.T) {
+	// Generate a new ECDSA private key
+	privKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("Failed to generate private key: %v", err)
+	}
+
+	// Encode the public key in uncompressed format
+	pubKeyBytes := elliptic.Marshal(crypto.S256(), privKey.PublicKey.X, privKey.PublicKey.Y)
+	pubKeyHex := hex.EncodeToString(pubKeyBytes)
+
+	// Calculate the expected address using common.BytesToAddress
+	hash := crypto.Keccak256(pubKeyBytes[1:])
+	expectedAddress := common.BytesToAddress(hash[12:]).String()
+
+	t.Logf("Public Key: %s", pubKeyHex)
+	t.Logf("Expected Address: %s", expectedAddress)
+
+	// Call the function PublicKeyToAddress
+	gotAddress, err := PublicKeyToAddress(pubKeyHex)
+	if err != nil {
+		t.Fatalf("Error creating address from public key: %v", err)
+	}
+
+	t.Logf("Got Address: %s", gotAddress)
+
+	// Compare the expected address with the address generated by the function
+	if gotAddress != expectedAddress {
+		t.Errorf("Expected address: %s, got: %s", expectedAddress, gotAddress)
+	}
+}