You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
360 lines
13 KiB
Python
360 lines
13 KiB
Python
"""
|
|
Tests for the Attention subsystem.
|
|
|
|
Covers:
|
|
- Entity resolver name normalization
|
|
- Wikipedia scoring logic
|
|
- GDELT query building
|
|
- Feature materializer calculations
|
|
- API endpoint routing and response schemas
|
|
"""
|
|
|
|
import pytest
|
|
from datetime import date
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from app.services.attention.entity_resolver import (
|
|
_normalize_name,
|
|
_score_wiki_result,
|
|
_build_gdelt_query,
|
|
_is_placeholder_name,
|
|
_fetch_sec_company_name,
|
|
)
|
|
from app.schemas.attention import (
|
|
EntityInfo,
|
|
WikiFeatures,
|
|
NewsFeatures,
|
|
EventAttentionResponse,
|
|
EntityResolveResponse,
|
|
CollectionStatusResponse,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entity Resolver: _normalize_name
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestNormalizeName:
|
|
def test_strips_inc(self):
|
|
canonical, aliases = _normalize_name("Apple Inc.")
|
|
assert canonical == "Apple"
|
|
assert "Apple Inc." in aliases
|
|
|
|
def test_strips_corp(self):
|
|
canonical, aliases = _normalize_name("Microsoft Corp")
|
|
assert canonical == "Microsoft"
|
|
|
|
def test_strips_holdings(self):
|
|
canonical, aliases = _normalize_name("Ondas Holdings Inc.")
|
|
assert canonical == "Ondas"
|
|
assert "Ondas Holdings" in aliases or "Ondas Holdings Inc." in aliases
|
|
|
|
def test_strips_multiple_suffixes(self):
|
|
canonical, aliases = _normalize_name("SomeCompany Holdings Ltd.")
|
|
assert canonical == "SomeCompany"
|
|
|
|
def test_no_suffix(self):
|
|
canonical, aliases = _normalize_name("Tesla")
|
|
assert canonical == "Tesla"
|
|
assert aliases == []
|
|
|
|
def test_preserves_original_in_aliases(self):
|
|
canonical, aliases = _normalize_name("Alphabet Inc.")
|
|
assert "Alphabet Inc." in aliases
|
|
|
|
def test_strips_technologies(self):
|
|
canonical, aliases = _normalize_name("Palantir Technologies Inc.")
|
|
assert canonical == "Palantir"
|
|
|
|
def test_strips_pharmaceuticals(self):
|
|
canonical, aliases = _normalize_name("Pfizer Pharmaceuticals Inc.")
|
|
assert canonical == "Pfizer"
|
|
|
|
def test_comma_handling(self):
|
|
canonical, aliases = _normalize_name("Berkshire Hathaway, Inc.")
|
|
assert canonical == "Berkshire Hathaway"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entity Resolver: _is_placeholder_name
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestIsPlaceholderName:
|
|
def test_ticker_equals_canonical(self):
|
|
assert _is_placeholder_name("AMZN", "AMZN") is True
|
|
|
|
def test_case_insensitive(self):
|
|
assert _is_placeholder_name("GOOGL", "googl") is True
|
|
|
|
def test_real_name_not_placeholder(self):
|
|
assert _is_placeholder_name("AAPL", "Apple") is False
|
|
|
|
def test_partial_ticker_not_placeholder(self):
|
|
assert _is_placeholder_name("META", "Meta Platforms") is False
|
|
|
|
def test_empty_canonical(self):
|
|
assert _is_placeholder_name("TSLA", "") is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entity Resolver: _fetch_sec_company_name
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestFetchSecCompanyName:
|
|
@pytest.mark.asyncio
|
|
async def test_returns_name_for_known_ticker(self):
|
|
import app.services.attention.entity_resolver as er
|
|
mock_data = {
|
|
"0": {"cik_str": 1018724, "ticker": "AMZN", "title": "AMAZON COM INC"},
|
|
"1": {"cik_str": 320193, "ticker": "AAPL", "title": "Apple Inc."},
|
|
}
|
|
with patch.object(er, "_SEC_TICKER_MAP", {}):
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_resp = MagicMock()
|
|
mock_resp.raise_for_status = MagicMock()
|
|
mock_resp.json.return_value = mock_data
|
|
mock_client_cls.return_value.__aenter__ = AsyncMock(
|
|
return_value=MagicMock(get=AsyncMock(return_value=mock_resp))
|
|
)
|
|
mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
result = await er._fetch_sec_company_name("AMZN")
|
|
assert result == "Amazon Com Inc"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_none_for_unknown_ticker(self):
|
|
import app.services.attention.entity_resolver as er
|
|
# Pre-populate cache with known tickers only
|
|
with patch.object(er, "_SEC_TICKER_MAP", {"AAPL": "Apple Inc."}):
|
|
result = await er._fetch_sec_company_name("ZZZZZ")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_cache_on_second_call(self):
|
|
import app.services.attention.entity_resolver as er
|
|
with patch.object(er, "_SEC_TICKER_MAP", {"NVDA": "NVIDIA CORP"}):
|
|
result = await er._fetch_sec_company_name("NVDA")
|
|
assert result == "Nvidia Corp"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_none_on_http_error(self):
|
|
import app.services.attention.entity_resolver as er
|
|
with patch.object(er, "_SEC_TICKER_MAP", {}):
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client_cls.return_value.__aenter__ = AsyncMock(
|
|
return_value=MagicMock(
|
|
get=AsyncMock(side_effect=Exception("network error"))
|
|
)
|
|
)
|
|
mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
result = await er._fetch_sec_company_name("AMZN")
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entity Resolver: _score_wiki_result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestScoreWikiResult:
|
|
def test_perfect_match_company_keyword(self):
|
|
result = {
|
|
"title": "Apple Inc.",
|
|
"snippet": "American multinational technology corporation and company stock",
|
|
}
|
|
score = _score_wiki_result(result, "Apple", [])
|
|
assert score > 0.5
|
|
|
|
def test_no_name_in_title_penalized(self):
|
|
result = {
|
|
"title": "Some Random Album",
|
|
"snippet": "Music album released in 2005",
|
|
}
|
|
score = _score_wiki_result(result, "Apple", [])
|
|
# Returns 0.1 (penalized for album signal), not 0.0
|
|
assert score <= 0.15
|
|
|
|
def test_penalizes_album(self):
|
|
result = {
|
|
"title": "Ondas (album)",
|
|
"snippet": "Ondas is a music album by some band",
|
|
}
|
|
score = _score_wiki_result(result, "Ondas", ["Ondas Holdings"])
|
|
assert score <= 0.15
|
|
|
|
def test_alias_match_in_title(self):
|
|
result = {
|
|
"title": "Ondas Holdings",
|
|
"snippet": "American company stock nasdaq finance",
|
|
}
|
|
score = _score_wiki_result(result, "Ondas", ["Ondas Holdings"])
|
|
assert score > 0.5
|
|
|
|
def test_film_penalized(self):
|
|
result = {
|
|
"title": "Terns (film)",
|
|
"snippet": "A 2003 film directed by someone",
|
|
}
|
|
score = _score_wiki_result(result, "Terns", [])
|
|
assert score <= 0.15
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entity Resolver: _build_gdelt_query
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestBuildGdeltQuery:
|
|
def test_single_term(self):
|
|
query = _build_gdelt_query("Apple", [])
|
|
assert query == '"Apple"'
|
|
|
|
def test_multiple_terms_joined_with_or(self):
|
|
query = _build_gdelt_query("Apple", ["Apple Inc."])
|
|
assert '"Apple"' in query
|
|
assert '"Apple Inc."' in query
|
|
assert " OR " in query
|
|
|
|
def test_caps_at_four_terms(self):
|
|
query = _build_gdelt_query("Tesla", ["Tesla Inc.", "Tesla Motors", "Tesla Corp", "Extra"])
|
|
parts = query.split(" OR ")
|
|
assert len(parts) == 4
|
|
|
|
def test_no_duplicate_canonical(self):
|
|
query = _build_gdelt_query("Palantir", ["Palantir Technologies Inc.", "Palantir Technologies"])
|
|
# canonical should appear only once
|
|
assert query.count('"Palantir"') == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schemas
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestSchemas:
|
|
def test_entity_info_defaults(self):
|
|
info = EntityInfo(ticker="AAPL", canonical_name="Apple")
|
|
assert info.aliases == []
|
|
assert info.resolver_confidence == 0.0
|
|
assert info.is_manual_override is False
|
|
assert info.wiki_title is None
|
|
|
|
def test_wiki_features_all_none(self):
|
|
wf = WikiFeatures()
|
|
assert wf.views is None
|
|
assert wf.spike_10d is None
|
|
assert wf.zscore_20d is None
|
|
|
|
def test_news_features_defaults(self):
|
|
nf = NewsFeatures()
|
|
assert nf.article_count_1d == 0
|
|
assert nf.article_count_3d == 0
|
|
assert nf.unique_domains_3d == 0
|
|
assert nf.us_article_count_3d == 0
|
|
assert nf.gdelt_status == "not_collected"
|
|
|
|
def test_news_features_gdelt_status_values(self):
|
|
assert NewsFeatures(gdelt_status="collected").gdelt_status == "collected"
|
|
assert NewsFeatures(gdelt_status="not_available").gdelt_status == "not_available"
|
|
|
|
def test_event_attention_response(self):
|
|
resp = EventAttentionResponse(
|
|
ticker="AAPL",
|
|
event_date=date(2026, 2, 6),
|
|
entity=EntityInfo(ticker="AAPL", canonical_name="Apple"),
|
|
wiki=WikiFeatures(views=10000, spike_10d=2.5, zscore_20d=1.8),
|
|
news=NewsFeatures(article_count_1d=5, article_count_3d=12),
|
|
)
|
|
assert resp.ticker == "AAPL"
|
|
assert resp.wiki.spike_10d == 2.5
|
|
assert resp.news.article_count_3d == 12
|
|
|
|
def test_collection_status_response(self):
|
|
resp = CollectionStatusResponse(
|
|
ticker="ONDS",
|
|
source="wiki",
|
|
records_collected=22,
|
|
date_range={"event_date": "2026-02-11"},
|
|
status="success",
|
|
)
|
|
assert resp.records_collected == 22
|
|
assert resp.source == "wiki"
|
|
|
|
def test_entity_resolve_response(self):
|
|
resp = EntityResolveResponse(
|
|
ticker="ONDS",
|
|
entity=EntityInfo(
|
|
ticker="ONDS",
|
|
canonical_name="Ondas",
|
|
wiki_title="Ondas Holdings",
|
|
resolver_confidence=0.75,
|
|
),
|
|
status="resolved",
|
|
message="Entity resolved",
|
|
)
|
|
assert resp.status == "resolved"
|
|
assert resp.entity.wiki_title == "Ondas Holdings"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Feature Materializer: stats calculations (unit-tested inline)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestStatCalculations:
|
|
"""Test the statistical computation logic without DB."""
|
|
|
|
def test_spike_calculation(self):
|
|
import statistics
|
|
wiki_views = 15000
|
|
prev_10 = [5000, 6000, 4500, 5500, 7000, 6500, 4000, 5000, 6000, 5500]
|
|
median_10 = statistics.median(prev_10)
|
|
spike = wiki_views / median_10
|
|
assert spike == pytest.approx(15000 / 5500, rel=1e-3)
|
|
|
|
def test_zscore_calculation(self):
|
|
import statistics
|
|
wiki_views = 15000
|
|
prev_20 = [5000] * 20
|
|
mean_20 = statistics.mean(prev_20)
|
|
stdev_20 = statistics.stdev(prev_20)
|
|
# All same values → stdev=0, no zscore
|
|
assert stdev_20 == 0.0
|
|
|
|
def test_zscore_with_variance(self):
|
|
import statistics
|
|
wiki_views = 15000
|
|
prev_20 = [4000, 5000, 6000, 4500, 5500, 7000, 3500, 4000, 5000, 6000,
|
|
4000, 5000, 6000, 4500, 5500, 7000, 3500, 4000, 5000, 6000]
|
|
mean_20 = statistics.mean(prev_20)
|
|
stdev_20 = statistics.stdev(prev_20)
|
|
zscore = (wiki_views - mean_20) / stdev_20
|
|
assert zscore > 0 # 15000 is well above mean ~5000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API endpoint integration (mock DB)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestAttentionEndpoints:
|
|
"""Lightweight route-level tests using FastAPI TestClient with mocked DB."""
|
|
|
|
@pytest.fixture
|
|
def client(self):
|
|
from fastapi import FastAPI
|
|
from app.api.v1.endpoints.attention import router
|
|
app = FastAPI()
|
|
app.include_router(router, prefix="/attention")
|
|
|
|
from fastapi.testclient import TestClient
|
|
return TestClient(app)
|
|
|
|
def test_entity_not_found_returns_404(self, client):
|
|
# Full integration tests require a live DB — this is a placeholder.
|
|
# Verified manually via: curl -X POST http://localhost:18001/api/v1/attention/admin/resolve/AAPL
|
|
pass
|
|
|
|
def test_routes_registered(self):
|
|
from app.api.v1.endpoints.attention import router
|
|
paths = [r.path for r in router.routes]
|
|
assert "/event/{ticker}" in paths
|
|
assert "/entity/{ticker}" in paths
|
|
assert "/admin/resolve/{ticker}" in paths
|
|
assert "/admin/collect/wiki/{ticker}" in paths
|
|
assert "/admin/collect/gdelt/{ticker}" in paths
|