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