204cd0e75b
- PR1: Project scaffold, DB models, price loader - PR2: Congressional trade ingestion (House Stock Watcher) - PR3: Security enrichment + deployment infrastructure - 37 passing tests, 87%+ coverage - Docker + Proxmox deployment ready - Complete documentation - Works 100% offline with fixtures
243 lines
7.4 KiB
Python
243 lines
7.4 KiB
Python
"""
|
|
Tests for security enricher.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sqlalchemy import select
|
|
|
|
from pote.db.models import Security
|
|
from pote.ingestion.security_enricher import SecurityEnricher
|
|
|
|
|
|
def test_enrich_security_success(test_db_session):
|
|
"""Test successful security enrichment."""
|
|
# Create an unenriched security (name == ticker)
|
|
security = Security(ticker="TSLA", name="TSLA", asset_type="stock")
|
|
test_db_session.add(security)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
# Mock yfinance response
|
|
mock_info = {
|
|
"symbol": "TSLA",
|
|
"longName": "Tesla, Inc.",
|
|
"sector": "Consumer Cyclical",
|
|
"industry": "Auto Manufacturers",
|
|
"exchange": "NASDAQ",
|
|
"quoteType": "EQUITY",
|
|
}
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = mock_info
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
success = enricher.enrich_security(security)
|
|
|
|
assert success is True
|
|
assert security.name == "Tesla, Inc."
|
|
assert security.sector == "Consumer Cyclical"
|
|
assert security.industry == "Auto Manufacturers"
|
|
assert security.exchange == "NASDAQ"
|
|
assert security.asset_type == "stock"
|
|
|
|
|
|
def test_enrich_security_etf(test_db_session):
|
|
"""Test enriching an ETF."""
|
|
security = Security(ticker="SPY", name="SPY", asset_type="stock")
|
|
test_db_session.add(security)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
mock_info = {
|
|
"symbol": "SPY",
|
|
"longName": "SPDR S&P 500 ETF Trust",
|
|
"sector": None,
|
|
"industry": None,
|
|
"exchange": "NYSE",
|
|
"quoteType": "ETF",
|
|
}
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = mock_info
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
success = enricher.enrich_security(security)
|
|
|
|
assert success is True
|
|
assert security.name == "SPDR S&P 500 ETF Trust"
|
|
assert security.asset_type == "etf"
|
|
|
|
|
|
def test_enrich_security_skip_already_enriched(test_db_session):
|
|
"""Test that already enriched securities are skipped by default."""
|
|
security = Security(
|
|
ticker="MSFT",
|
|
name="Microsoft Corporation", # Already enriched
|
|
sector="Technology",
|
|
asset_type="stock",
|
|
)
|
|
test_db_session.add(security)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
# Should skip without calling yfinance
|
|
success = enricher.enrich_security(security, force=False)
|
|
assert success is False
|
|
|
|
|
|
def test_enrich_security_force_refresh(test_db_session):
|
|
"""Test force re-enrichment."""
|
|
security = Security(
|
|
ticker="GOOGL",
|
|
name="Alphabet Inc.", # Already enriched
|
|
sector="Technology",
|
|
asset_type="stock",
|
|
)
|
|
test_db_session.add(security)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
mock_info = {
|
|
"symbol": "GOOGL",
|
|
"longName": "Alphabet Inc. Class A", # Updated name
|
|
"sector": "Communication Services", # Updated sector
|
|
"industry": "Internet Content & Information",
|
|
"exchange": "NASDAQ",
|
|
"quoteType": "EQUITY",
|
|
}
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = mock_info
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
success = enricher.enrich_security(security, force=True)
|
|
|
|
assert success is True
|
|
assert security.name == "Alphabet Inc. Class A"
|
|
assert security.sector == "Communication Services"
|
|
|
|
|
|
def test_enrich_security_no_data(test_db_session, sample_security):
|
|
"""Test handling of ticker with no data."""
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
# Mock empty response
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = {} # No data
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
success = enricher.enrich_security(sample_security)
|
|
|
|
assert success is False
|
|
# Original values should be unchanged
|
|
assert sample_security.name == "Apple Inc."
|
|
|
|
|
|
def test_enrich_all_securities(test_db_session):
|
|
"""Test enriching multiple securities."""
|
|
# Create unenriched securities (name == ticker)
|
|
securities = [
|
|
Security(ticker="AAPL", name="AAPL", asset_type="stock"),
|
|
Security(ticker="MSFT", name="MSFT", asset_type="stock"),
|
|
Security(ticker="GOOGL", name="GOOGL", asset_type="stock"),
|
|
]
|
|
for sec in securities:
|
|
test_db_session.add(sec)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
def mock_info_fn(ticker):
|
|
return {
|
|
"symbol": ticker,
|
|
"longName": f"{ticker} Corporation",
|
|
"sector": "Technology",
|
|
"industry": "Software",
|
|
"exchange": "NASDAQ",
|
|
"quoteType": "EQUITY",
|
|
}
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
|
|
def side_effect(ticker_str):
|
|
mock_instance = MagicMock()
|
|
mock_instance.info = mock_info_fn(ticker_str)
|
|
return mock_instance
|
|
|
|
mock_ticker.side_effect = side_effect
|
|
|
|
counts = enricher.enrich_all_securities()
|
|
|
|
assert counts["total"] == 3
|
|
assert counts["enriched"] == 3
|
|
assert counts["failed"] == 0
|
|
|
|
# Verify enrichment
|
|
stmt = select(Security).where(Security.ticker == "AAPL")
|
|
aapl = test_db_session.scalars(stmt).first()
|
|
assert aapl.name == "AAPL Corporation"
|
|
assert aapl.sector == "Technology"
|
|
|
|
|
|
def test_enrich_all_securities_with_limit(test_db_session):
|
|
"""Test enriching with a limit."""
|
|
# Create 5 unenriched securities
|
|
for i in range(5):
|
|
security = Security(ticker=f"TEST{i}", name=f"TEST{i}", asset_type="stock")
|
|
test_db_session.add(security)
|
|
test_db_session.commit()
|
|
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = {
|
|
"symbol": "TEST",
|
|
"longName": "Test Corp",
|
|
"quoteType": "EQUITY",
|
|
}
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
counts = enricher.enrich_all_securities(limit=2)
|
|
|
|
assert counts["total"] == 2
|
|
assert counts["enriched"] == 2
|
|
|
|
|
|
def test_enrich_by_ticker_success(test_db_session, sample_security):
|
|
"""Test enriching by specific ticker."""
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
mock_info = {
|
|
"symbol": "AAPL",
|
|
"longName": "Apple Inc.",
|
|
"sector": "Technology",
|
|
"quoteType": "EQUITY",
|
|
}
|
|
|
|
with patch("pote.ingestion.security_enricher.yf.Ticker") as mock_ticker:
|
|
mock_ticker_instance = MagicMock()
|
|
mock_ticker_instance.info = mock_info
|
|
mock_ticker.return_value = mock_ticker_instance
|
|
|
|
success = enricher.enrich_by_ticker("AAPL")
|
|
|
|
assert success is True
|
|
|
|
|
|
def test_enrich_by_ticker_not_found(test_db_session):
|
|
"""Test enriching a ticker not in database."""
|
|
enricher = SecurityEnricher(test_db_session)
|
|
|
|
success = enricher.enrich_by_ticker("NOTFOUND")
|
|
assert success is False
|