Fix NFT mock detector test and ensure proper initialization

Author: Jetlum

mock/nft_mock.go

@@ -14,16 +14,13 @@ type MockNFTDetector struct {
 }
 
 func NewMockNFTDetector() *MockNFTDetector {
-	detector := &MockNFTDetector{
-		IsNFTTransactionFunc: func(tx *types.Transaction) bool {
-			return false
-		},
-	}
+	detector := &MockNFTDetector{}
+	detector.knownContracts = sync.Map{}
 
-	// Initialize with some test NFT contract addresses
+	// Initialize with test contracts
 	testContracts := map[string]bool{
-		"0xBC4CA0EdA7647A8aB7C2061c2E118A18a936f13D": true, // Test BAYC
-		"0x23581767a106ae21c074b2276D25e5C3e136a68b": true, // Test Moonbirds
+		"0xBC4CA0EdA7647A8aB7C2061c2E118A18a936f13D": true,
+		"0x23581767a106ae21c074b2276D25e5C3e136a68b": true,
 	}
 
 	for addr := range testContracts {
@@ -36,6 +33,7 @@ func NewMockNFTDetector() *MockNFTDetector {
 func (m *MockNFTDetector) IsNFTTransaction(tx *types.Transaction) bool {
 	m.callCount++
 
+	// Custom function takes precedence
 	if m.IsNFTTransactionFunc != nil {
 		return m.IsNFTTransactionFunc(tx)
 	}
@@ -44,11 +42,13 @@ func (m *MockNFTDetector) IsNFTTransaction(tx *types.Transaction) bool {
 		return false
 	}
 
-	if val, exists := m.knownContracts.Load(*tx.To()); exists {
-		return val.(bool)
+	toAddr := tx.To()
+	// Check known contracts
+	val, exists := m.knownContracts.Load(*toAddr)
+	if !exists {
+		return false
 	}
-
-	return false
+	return val.(bool)
 }
 
 // Helper methods for testing

mock/nft_mock_test.go

@@ -9,35 +9,28 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+// mock/nft_mock_test.go
 func TestMockNFTDetector(t *testing.T) {
-	detector := NewMockNFTDetector()
-
 	t.Run("Known contract", func(t *testing.T) {
-		detector.Reset()
-		address := common.HexToAddress("0xBAYC")
-		detector.AddKnownContract(address)
+		detector := NewMockNFTDetector() // Create fresh instance
+
+		contractAddr := common.HexToAddress("0xBC4CA0EdA7647A8aB7C2061c2E118A18a936f13D")
+
+		// Verify initial state
+		val, exists := detector.knownContracts.Load(contractAddr)
+		assert.True(t, exists, "Contract should be in initial known contracts")
+		assert.True(t, val.(bool), "Contract should be marked as valid")
 
-		// Create transaction to the known contract address
 		tx := types.NewTransaction(
 			0,
-			address,
+			contractAddr,
 			big.NewInt(0),
 			21000,
 			big.NewInt(1),
 			nil,
 		)
 
-		// Add debug assertions
-		toAddr := tx.To()
-		if toAddr == nil {
-			t.Fatal("Transaction To address is nil")
-		}
-
-		// Verify the addresses match
-		assert.Equal(t, address, *toAddr, "Transaction address should match known contract")
-
-		// Check if transaction is detected as NFT
 		result := detector.IsNFTTransaction(tx)
-		assert.True(t, result, "Transaction to known NFT contract address should return true")
+		assert.True(t, result, "Transaction to known NFT contract should be detected")
 	})
 }