185 lines
6 KiB
Python
185 lines
6 KiB
Python
"""
|
|
Tests for Wallet Drain Scanner
|
|
=================================
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
from app.wallet_drain_scanner import (
|
|
WalletDrainScanner,
|
|
_classify_drain_risk,
|
|
_compute_drain_score,
|
|
_generate_recommendation,
|
|
)
|
|
|
|
|
|
class TestDrainScore:
|
|
"""Test the core scoring algorithm."""
|
|
|
|
def test_zero_score(self):
|
|
"""No signals should give near-zero score."""
|
|
score = _compute_drain_score(
|
|
unlimited_approvals=0,
|
|
known_drainer_txs=0,
|
|
suspicious_permit_count=0,
|
|
max_approval_amount_score=0.0,
|
|
nft_unlimited_approvals=0,
|
|
delegate_call_count=0,
|
|
)
|
|
assert 0 <= score <= 5
|
|
|
|
def test_critical_drain(self):
|
|
"""All drain signals maxed should give very high score."""
|
|
score = _compute_drain_score(
|
|
unlimited_approvals=4,
|
|
known_drainer_txs=4,
|
|
suspicious_permit_count=4,
|
|
max_approval_amount_score=1.0,
|
|
nft_unlimited_approvals=4,
|
|
delegate_call_count=4,
|
|
)
|
|
assert score >= 75
|
|
|
|
def test_moderate_drain(self):
|
|
"""Mixed signals should give moderate score."""
|
|
score = _compute_drain_score(
|
|
unlimited_approvals=1,
|
|
known_drainer_txs=1,
|
|
suspicious_permit_count=1,
|
|
max_approval_amount_score=0.5,
|
|
nft_unlimited_approvals=0,
|
|
delegate_call_count=1,
|
|
)
|
|
assert 20 <= score <= 75
|
|
|
|
def test_score_boundary(self):
|
|
"""Score should never exceed 100 or go below 0."""
|
|
score = _compute_drain_score(
|
|
unlimited_approvals=999,
|
|
known_drainer_txs=999,
|
|
suspicious_permit_count=999,
|
|
max_approval_amount_score=999.0,
|
|
nft_unlimited_approvals=999,
|
|
delegate_call_count=999,
|
|
)
|
|
assert score <= 100.0
|
|
assert score >= 0.0
|
|
|
|
def test_unlimited_approval_weight(self):
|
|
"""More unlimited approvals should increase score."""
|
|
low = _compute_drain_score(
|
|
unlimited_approvals=0,
|
|
known_drainer_txs=0,
|
|
suspicious_permit_count=0,
|
|
max_approval_amount_score=0.0,
|
|
nft_unlimited_approvals=0,
|
|
delegate_call_count=0,
|
|
)
|
|
high = _compute_drain_score(
|
|
unlimited_approvals=3,
|
|
known_drainer_txs=0,
|
|
suspicious_permit_count=0,
|
|
max_approval_amount_score=0.0,
|
|
nft_unlimited_approvals=0,
|
|
delegate_call_count=0,
|
|
)
|
|
assert high > low
|
|
|
|
|
|
class TestClassification:
|
|
"""Test drain risk classification thresholds."""
|
|
|
|
def test_critical(self):
|
|
assert _classify_drain_risk(80) == "critical"
|
|
assert _classify_drain_risk(70) == "critical"
|
|
|
|
def test_high(self):
|
|
assert _classify_drain_risk(65) == "high"
|
|
assert _classify_drain_risk(50) == "high"
|
|
|
|
def test_moderate(self):
|
|
assert _classify_drain_risk(45) == "moderate"
|
|
assert _classify_drain_risk(30) == "moderate"
|
|
|
|
def test_low(self):
|
|
assert _classify_drain_risk(25) == "low"
|
|
assert _classify_drain_risk(10) == "low"
|
|
|
|
def test_none(self):
|
|
assert _classify_drain_risk(5) == "none"
|
|
assert _classify_drain_risk(0) == "none"
|
|
|
|
|
|
class TestRecommendation:
|
|
"""Test recommendation generation."""
|
|
|
|
def test_critical_recommendation(self):
|
|
rec = _generate_recommendation(85, 0.9)
|
|
assert "CRITICAL" in rec
|
|
assert "revoke" in rec.lower()
|
|
|
|
def test_high_recommendation(self):
|
|
rec = _generate_recommendation(60, 0.5)
|
|
assert "HIGH" in rec
|
|
|
|
def test_moderate_recommendation(self):
|
|
rec = _generate_recommendation(40, 0.3)
|
|
assert "MODERATE" in rec or "risky" in rec
|
|
|
|
def test_low_with_dangerous_ratio(self):
|
|
rec = _generate_recommendation(15, 0.1)
|
|
assert "LOW" in rec or "Minor" in rec
|
|
|
|
def test_none_recommendation(self):
|
|
rec = _generate_recommendation(5, 0.0)
|
|
assert "No drain" in rec or "healthy" in rec
|
|
|
|
def test_different_severity_format(self):
|
|
"""Different severities should produce different messages."""
|
|
low_rec = _generate_recommendation(5, 0.0)
|
|
high_rec = _generate_recommendation(80, 0.9)
|
|
assert low_rec != high_rec
|
|
assert "CRITICAL" in high_rec
|
|
|
|
|
|
class TestWalletDrainScanner:
|
|
"""Test the WalletDrainScanner class."""
|
|
|
|
def test_invalid_address_raises(self):
|
|
"""Invalid address should raise ValueError."""
|
|
with pytest.raises(ValueError, match="Invalid address"):
|
|
WalletDrainScanner("nope", "ethereum")
|
|
|
|
def test_valid_evm_address(self):
|
|
"""Valid EVM address should create scanner."""
|
|
scanner = WalletDrainScanner(
|
|
"0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
|
|
"ethereum",
|
|
)
|
|
assert scanner.address == "0x7a250d5630b4cf539739df2c5dacb4c659f2488d"
|
|
assert scanner.chain == "ethereum"
|
|
assert scanner.is_evm
|
|
|
|
def test_known_drainer_detection_bad_prefix(self):
|
|
"""Addresses with 0x0000 prefix should be detected as drainer."""
|
|
assert WalletDrainScanner._is_known_drainer("0x000000000000000000000000000000000000dead")
|
|
|
|
def test_known_drainer_benign_not_detected(self):
|
|
"""Well-known benign contracts should NOT be detected as drainers."""
|
|
assert not WalletDrainScanner._is_known_drainer("0x7a250d5630b4cf539739df2c5dacb4c659f2488d")
|
|
|
|
def test_known_drainer_dead_prefix(self):
|
|
"""Addresses with 0xdead prefix should be flagged."""
|
|
assert WalletDrainScanner._is_known_drainer("0xdead000000000000000000000000000000000000")
|
|
|
|
def test_validate_url_rejects_invalid(self):
|
|
"""Malformed URLs should be rejected."""
|
|
from app.wallet_drain_scanner import _validate_url
|
|
|
|
assert not _validate_url("javascript:alert(1)")
|
|
assert _validate_url("https://api.etherscan.io/api")
|