backend/app/test_arkham_entity.py

"""
Tests for Arkham Entity Resolver (arkham_entity.py)
"""

import asyncio
import json
import unittest
from datetime import UTC, datetime

from app.arkham_entity import (
    ArkhamEntityResolver,
    EntityCategory,
    EntityLabel,
    EntityReport,
    ResolvedEntity,
    ResolverSource,
    analyze_addresses,
    format_report,
)


class TestEntityModels(unittest.TestCase):
    """Test dataclass construction and serialization."""

    def test_entity_label_defaults(self):
        """EntityLabel should accept all fields."""
        label = EntityLabel(
            label="Binance Hot Wallet",
            source=ResolverSource.LOCAL_DB,
            confidence=0.95,
            category="exchange",
        )
        self.assertEqual(label.label, "Binance Hot Wallet")
        self.assertEqual(label.source, ResolverSource.LOCAL_DB)
        self.assertEqual(label.confidence, 0.95)
        self.assertEqual(label.category, "exchange")

    def test_resolved_entity_defaults(self):
        """ResolvedEntity should set default resolved_at."""
        entity = ResolvedEntity(
            address="0xBE0eB53FC46b790099138e3d32C721856d41e865",
            chain="ethereum",
            entity_name="Binance",
            category=EntityCategory.EXCHANGE,
            confidence=95.0,
        )
        self.assertEqual(entity.entity_name, "Binance")
        self.assertEqual(entity.category, EntityCategory.EXCHANGE)
        self.assertEqual(entity.confidence, 95.0)
        self.assertIsNotNone(entity.resolved_at)
        self.assertEqual(len(entity.labels), 0)
        self.assertEqual(len(entity.tags), 0)

    def test_resolved_entity_to_dict(self):
        """to_dict should produce a clean serializable dict."""
        entity = ResolvedEntity(
            address="0xBE0eB53FC46b790099138e3d32C721856d41e865",
            chain="ethereum",
            entity_name="Binance",
            category=EntityCategory.EXCHANGE,
            confidence=95.0,
            labels=[
                EntityLabel(
                    label="Binance Hot Wallet 7",
                    source=ResolverSource.LOCAL_DB,
                    confidence=0.95,
                    category="exchange",
                )
            ],
            tags=["cex", "hot_wallet"],
        )
        d = entity.to_dict()
        self.assertEqual(d["entity_name"], "Binance")
        self.assertEqual(d["category"], "exchange")
        self.assertEqual(d["confidence"], 95.0)
        self.assertEqual(len(d["labels"]), 1)
        self.assertEqual(d["labels"][0]["label"], "Binance Hot Wallet 7")

    def test_entity_report_defaults(self):
        """EntityReport should handle empty entities."""
        report = EntityReport(
            query_addresses=["0x1234"],
            chain="ethereum",
            entities=[],
            summary={"total": 0, "known": 0, "unknown": 0, "avg_confidence": 0.0},
        )
        self.assertEqual(len(report.entities), 0)
        self.assertIsNone(report.error)

    def test_entity_report_error(self):
        """EntityReport should carry error state."""
        report = EntityReport(
            query_addresses=[],
            chain="unknown",
            entities=[],
            summary={"total": 0},
            error="No addresses provided",
        )
        self.assertEqual(report.error, "No addresses provided")


class TestEntityCategory(unittest.TestCase):
    """Test EntityCategory enum."""

    def test_all_categories_exist(self):
        """All expected entity categories should be defined."""
        expected = [
            "exchange", "defi", "token", "bridge", "fund",
            "whale", "nft", "gaming", "scam", "sanctioned", "unknown",
        ]
        for cat in expected:
            self.assertIn(cat, EntityCategory._value2member_map_)

    def test_category_values(self):
        """Each category should map to its string value."""
        self.assertEqual(EntityCategory.EXCHANGE.value, "exchange")
        self.assertEqual(EntityCategory.DEFI.value, "defi")
        self.assertEqual(EntityCategory.UNKNOWN.value, "unknown")
        self.assertEqual(EntityCategory.SCAM.value, "scam")


class TestResolverSource(unittest.TestCase):
    """Test ResolverSource enum."""

    def test_all_sources_exist(self):
        """All expected resolver sources should be defined."""
        expected = ["arkham_api", "local_db", "entity_registry", "entity_labeler", "heuristic"]
        for src in expected:
            self.assertIn(src, ResolverSource._value2member_map_)


class TestArkhamEntityResolver(unittest.TestCase):
    """Test suite for ArkhamEntityResolver."""

    def setUp(self):
        self.resolver = ArkhamEntityResolver()

    # ═══════════════════════════════════════════════════════════════════
    # Smoke Tests
    # ═══════════════════════════════════════════════════════════════════

    def test_invalid_address_returns_low_confidence(self):
        """Invalid address should return 0 confidence with 'Invalid Address'."""
        entity = asyncio.run(self.resolver.resolve("not_an_address"))
        self.assertEqual(entity.confidence, 0.0)
        self.assertEqual(entity.entity_name, "Invalid Address")
        self.assertEqual(entity.category, EntityCategory.UNKNOWN)

    def test_empty_address_returns_low_confidence(self):
        """Empty string should be treated as invalid."""
        entity = asyncio.run(self.resolver.resolve(""))
        self.assertEqual(entity.confidence, 0.0)

    # ═══════════════════════════════════════════════════════════════════
    # Local DB Lookup Tests
    # ═══════════════════════════════════════════════════════════════════

    def test_known_binance_address_resolved(self):
        """Binance hot wallet should be resolved from local DB."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0xBE0eB53FC46b790099138e3d32C721856d41e865"
            )
        )
        self.assertEqual(entity.entity_name, "Binance")
        self.assertEqual(entity.category, EntityCategory.EXCHANGE)
        self.assertGreaterEqual(entity.confidence, 90.0)
        self.assertIn("cex", entity.tags)
        self.assertIn("hot_wallet", entity.tags)
        self.assertEqual(entity.source, ResolverSource.LOCAL_DB)

    def test_known_coinbase_address_resolved(self):
        """Coinbase hot wallet should be resolved from local DB."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0x503828976D22510aad0201ac7EC88293211D23Da"
            )
        )
        self.assertEqual(entity.entity_name, "Coinbase")
        self.assertEqual(entity.category, EntityCategory.EXCHANGE)

    def test_known_uniswap_router_resolved(self):
        """Uniswap V2 Router should be resolved from local DB."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0x7a250d5630b4cf139281983dce37532e7d5c9196"
            )
        )
        self.assertEqual(entity.entity_name, "Uniswap V2 Router")
        self.assertEqual(entity.category, EntityCategory.DEFI)

    def test_known_tether_contract_resolved(self):
        """Tether USDT contract should be resolved from local DB."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0xdAC17F958D2ee523a2206206994597C13D831ec7"
            )
        )
        self.assertEqual(entity.entity_name, "Tether (USDT)")
        self.assertEqual(entity.category, EntityCategory.TOKEN)

    def test_lazarus_group_resolved(self):
        """Sanctioned entity should be resolved from local DB."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0x1CBd3b2770909D4e10f157cABC84C7264073C9Ec"
            )
        )
        self.assertIn("Lazarus", entity.entity_name)
        self.assertIn("sanctioned", entity.tags)

    # ═══════════════════════════════════════════════════════════════════
    # Unknown Address Tests
    # ═══════════════════════════════════════════════════════════════════

    def test_unknown_address_returns_unknown(self):
        """Unknown address should return 'Unknown' with low confidence."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0x1234567890abcdef1234567890abcdef12345678",
                use_arkham=False,
            )
        )
        # Should be unknown since it's not in any database
        self.assertEqual(entity.entity_name, "Unknown")
        self.assertEqual(entity.category, EntityCategory.UNKNOWN)
        self.assertEqual(entity.confidence, 0.0)

    # ═══════════════════════════════════════════════════════════════════
    # Chain Detection Tests
    # ═══════════════════════════════════════════════════════════════════

    def test_ethereum_address_detected(self):
        """0x-prefixed 40-hex-char addresses should detect as ethereum."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0xBE0eB53FC46b790099138e3d32C721856d41e865"
            )
        )
        self.assertEqual(entity.chain, "ethereum")

    def test_solana_address_detected(self):
        """Base58 addresses should detect as solana."""
        entity = asyncio.run(
            self.resolver.resolve(
                "9WzDXwBbmkg8ZTbNMqUxvQRAyrZzDsGYdLVL9zYtAWWM"
            )
        )
        self.assertEqual(entity.chain, "solana")

    def test_explicit_chain_override(self):
        """Explicit chain parameter should override detection."""
        entity = asyncio.run(
            self.resolver.resolve(
                "0xBE0eB53FC46b790099138e3d32C721856d41e865",
                chain="polygon",
            )
        )
        self.assertEqual(entity.chain, "polygon")


class TestAnalyzeAddresses(unittest.TestCase):
    """Test the high-level analyze_addresses function."""

    def test_empty_addresses_returns_error(self):
        """Empty address list should return an error report."""
        report = asyncio.run(analyze_addresses([]))
        self.assertIsNotNone(report.error)
        self.assertEqual(report.error, "No addresses provided")

    def test_single_known_address(self):
        """Single known address should produce a correct report."""
        report = asyncio.run(
            analyze_addresses(
                ["0xBE0eB53FC46b790099138e3d32C721856d41e865"],
                use_arkham=False,
            )
        )
        self.assertEqual(len(report.entities), 1)
        self.assertEqual(report.entities[0].entity_name, "Binance")
        self.assertIsNone(report.error)
        self.assertEqual(report.summary["total"], 1)
        self.assertEqual(report.summary["known"], 1)

    def test_multiple_addresses_mixed(self):
        """Mixed known/unknown addresses should both appear."""
        report = asyncio.run(
            analyze_addresses(
                [
                    "0xBE0eB53FC46b790099138e3d32C721856d41e865",  # Binance
                    "0x0000000000000000000000000000000000000001",  # Unknown
                ],
                use_arkham=False,
            )
        )
        self.assertEqual(len(report.entities), 2)
        self.assertEqual(report.entities[0].entity_name, "Binance")
        self.assertEqual(report.entities[1].entity_name, "Unknown")

    def test_report_summary(self):
        """Report summary should have correct aggregation."""
        report = asyncio.run(
            analyze_addresses(
                [
                    "0xBE0eB53FC46b790099138e3d32C721856d41e865",  # Binance
                    "0x503828976D22510aad0201ac7EC88293211D23Da",  # Coinbase
                ],
                use_arkham=False,
            )
        )
        self.assertEqual(report.summary["total"], 2)
        self.assertEqual(report.summary["known"], 2)
        self.assertEqual(report.summary["unknown"], 0)
        self.assertGreater(report.summary["avg_confidence"], 90.0)


class TestFormatReport(unittest.TestCase):
    """Test the format_report output."""

    def test_format_report_contains_entity_names(self):
        """Formatted report should include entity names."""
        report = asyncio.run(
            analyze_addresses(
                ["0xBE0eB53FC46b790099138e3d32C721856d41e865"],
                use_arkham=False,
            )
        )
        output = format_report(report)
        self.assertIn("Binance", output)
        self.assertIn("exchange", output)
        self.assertIn("confidence", output.lower())

    def test_format_report_error(self):
        """Error report should include the error message."""
        report = EntityReport(
            query_addresses=[],
            chain="unknown",
            entities=[],
            summary={"total": 0},
            error="No addresses provided",
        )
        output = format_report(report)
        self.assertIn("ERROR", output)
        self.assertIn("No addresses provided", output)

    def test_format_report_multiple_entities(self):
        """Report with multiple entities should show all."""
        report = asyncio.run(
            analyze_addresses(
                [
                    "0xBE0eB53FC46b790099138e3d32C721856d41e865",
                    "0x503828976D22510aad0201ac7EC88293211D23Da",
                ],
                use_arkham=False,
            )
        )
        output = format_report(report)
        self.assertIn("Binance", output)
        self.assertIn("Coinbase", output)


class TestBatchResolve(unittest.TestCase):
    """Test the batch resolve functionality."""

    def test_batch_resolve_empty(self):
        """Empty batch should return empty list."""
        entities = asyncio.run(self._empty_batch())
        self.assertEqual(len(entities), 0)

    def _empty_batch(self):
        r = ArkhamEntityResolver()
        return r.resolve_batch([], use_arkham=False)

    def test_batch_resolve_multiple(self):
        """Batch resolve should return results for all addresses."""
        async def _run():
            r = ArkhamEntityResolver()
            results = await r.resolve_batch(
                [
                    "0xBE0eB53FC46b790099138e3d32C721856d41e865",
                    "invalid_address",
                ],
                use_arkham=False,
            )
            await r.close()
            return results

        entities = asyncio.run(_run())
        self.assertEqual(len(entities), 2)
        self.assertEqual(entities[0].entity_name, "Binance")
        self.assertEqual(entities[1].entity_name, "Invalid Address")


if __name__ == "__main__":
    unittest.main()