83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
"""Tests for entity extraction"""
|
|
|
|
import pytest
|
|
from sentiment_engine.nlp.entity_extraction import AssetMapper, EntityExtractor
|
|
|
|
|
|
class TestAssetMapper:
|
|
"""Test asset mapping"""
|
|
|
|
def test_map_known_ticker(self):
|
|
mapper = AssetMapper()
|
|
asset_id, confidence = mapper.map_ticker("BTC")
|
|
assert asset_id == "BTC"
|
|
assert confidence >= 0.9
|
|
|
|
def test_map_alias(self):
|
|
mapper = AssetMapper()
|
|
asset_id, confidence = mapper.map_ticker("VITALIK")
|
|
assert asset_id == "ETH"
|
|
assert confidence >= 0.7
|
|
|
|
def test_map_unknown_ticker(self):
|
|
mapper = AssetMapper()
|
|
asset_id, confidence = mapper.map_ticker("UNKNOWNTICKER")
|
|
assert asset_id == "UNKNOWNTICKER"
|
|
assert confidence == 0.5
|
|
|
|
def test_map_contract(self):
|
|
mapper = AssetMapper()
|
|
asset_id, confidence, chain = mapper.map_contract("0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2")
|
|
assert asset_id == "ETH"
|
|
assert confidence >= 0.9
|
|
assert chain == "ethereum"
|
|
|
|
def test_resolve_aliases(self):
|
|
mapper = AssetMapper()
|
|
results = mapper.resolve_alias("Vitalik Buterin says ETH will moon")
|
|
assert any(r[1] == "ETH" for r in results)
|
|
|
|
|
|
class TestEntityExtractor:
|
|
"""Test entity extraction"""
|
|
|
|
@pytest.fixture
|
|
def extractor(self):
|
|
return EntityExtractor()
|
|
|
|
def test_extract_tickers(self, extractor):
|
|
text = "BTC and ETH are pumping hard"
|
|
mentions = extractor.extract_tickers(text)
|
|
assert len(mentions) == 2
|
|
asset_ids = [m.asset_id for m in mentions]
|
|
assert "BTC" in asset_ids
|
|
assert "ETH" in asset_ids
|
|
|
|
def test_extract_contracts(self, extractor):
|
|
text = "Send to 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
|
|
mentions = extractor.extract_contracts(text)
|
|
assert len(mentions) == 1
|
|
assert mentions[0].asset_id == "ETH"
|
|
|
|
def test_extract_aliases(self, extractor):
|
|
text = "Vitalik says ETH to the moon"
|
|
mentions = extractor.extract_aliases(text)
|
|
assert len(mentions) >= 1
|
|
assert mentions[0].asset_id == "ETH"
|
|
|
|
def test_deduplication(self, extractor):
|
|
text = "BTC BTC BTC"
|
|
mentions = extractor.extract_tickers(text)
|
|
assert len(mentions) == 1
|
|
assert mentions[0].asset_id == "BTC"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_all(self, extractor):
|
|
text = "BTC surges. Vitalik buys ETH. Send to 0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
|
|
entities = await extractor.extract_all(text)
|
|
asset_ids = [e.asset_id for e in entities]
|
|
assert "BTC" in asset_ids
|
|
assert "ETH" in asset_ids
|
|
# Contract should also map to ETH
|
|
assert asset_ids.count("ETH") >= 1
|