tests/unit/test_wallet_drain_scanner.py

"""
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")