Merge pull request 'feat: integrate TelegramClient into main trading loop (issue #34)' (#39) from feature/issue-34-main-integration into main
Some checks failed
CI / test (push) Has been cancelled

Reviewed-on: #39
This commit was merged in pull request #39.
This commit is contained in:
2026-02-04 23:44:49 +09:00
2 changed files with 425 additions and 7 deletions

View File

@@ -10,6 +10,7 @@ import argparse
import asyncio import asyncio
import logging import logging
import signal import signal
import sys
from datetime import UTC, datetime from datetime import UTC, datetime
from typing import Any from typing import Any
@@ -23,11 +24,12 @@ from src.context.layer import ContextLayer
from src.context.store import ContextStore from src.context.store import ContextStore
from src.core.criticality import CriticalityAssessor from src.core.criticality import CriticalityAssessor
from src.core.priority_queue import PriorityTaskQueue from src.core.priority_queue import PriorityTaskQueue
from src.core.risk_manager import CircuitBreakerTripped, RiskManager from src.core.risk_manager import CircuitBreakerTripped, FatFingerRejected, RiskManager
from src.db import init_db, log_trade from src.db import init_db, log_trade
from src.logging.decision_logger import DecisionLogger from src.logging.decision_logger import DecisionLogger
from src.logging_config import setup_logging from src.logging_config import setup_logging
from src.markets.schedule import MarketInfo, get_next_market_open, get_open_markets from src.markets.schedule import MarketInfo, get_next_market_open, get_open_markets
from src.notifications.telegram_client import TelegramClient
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -62,6 +64,7 @@ async def trading_cycle(
decision_logger: DecisionLogger, decision_logger: DecisionLogger,
context_store: ContextStore, context_store: ContextStore,
criticality_assessor: CriticalityAssessor, criticality_assessor: CriticalityAssessor,
telegram: TelegramClient,
market: MarketInfo, market: MarketInfo,
stock_code: str, stock_code: str,
) -> None: ) -> None:
@@ -199,11 +202,23 @@ async def trading_cycle(
order_amount = current_price * quantity order_amount = current_price * quantity
# 4. Risk check BEFORE order # 4. Risk check BEFORE order
try:
risk.validate_order( risk.validate_order(
current_pnl_pct=pnl_pct, current_pnl_pct=pnl_pct,
order_amount=order_amount, order_amount=order_amount,
total_cash=total_cash, total_cash=total_cash,
) )
except FatFingerRejected as exc:
try:
await telegram.notify_fat_finger(
stock_code=stock_code,
order_amount=exc.order_amount,
total_cash=exc.total_cash,
max_pct=exc.max_pct,
)
except Exception as notify_exc:
logger.warning("Fat finger notification failed: %s", notify_exc)
raise # Re-raise to prevent trade
# 5. Send order # 5. Send order
if market.is_domestic: if market.is_domestic:
@@ -223,6 +238,19 @@ async def trading_cycle(
) )
logger.info("Order result: %s", result.get("msg1", "OK")) logger.info("Order result: %s", result.get("msg1", "OK"))
# 5.5. Notify trade execution
try:
await telegram.notify_trade_execution(
stock_code=stock_code,
market=market.name,
action=decision.action,
quantity=quantity,
price=current_price,
confidence=decision.confidence,
)
except Exception as exc:
logger.warning("Telegram notification failed: %s", exc)
# 6. Log trade # 6. Log trade
log_trade( log_trade(
conn=db_conn, conn=db_conn,
@@ -266,6 +294,13 @@ async def run(settings: Settings) -> None:
decision_logger = DecisionLogger(db_conn) decision_logger = DecisionLogger(db_conn)
context_store = ContextStore(db_conn) context_store = ContextStore(db_conn)
# Initialize Telegram notifications
telegram = TelegramClient(
bot_token=settings.TELEGRAM_BOT_TOKEN,
chat_id=settings.TELEGRAM_CHAT_ID,
enabled=settings.TELEGRAM_ENABLED,
)
# Initialize volatility hunter # Initialize volatility hunter
volatility_analyzer = VolatilityAnalyzer(min_volume_surge=2.0, min_price_change=1.0) volatility_analyzer = VolatilityAnalyzer(min_volume_surge=2.0, min_price_change=1.0)
market_scanner = MarketScanner( market_scanner = MarketScanner(
@@ -289,6 +324,9 @@ async def run(settings: Settings) -> None:
# Track last scan time for each market # Track last scan time for each market
last_scan_time: dict[str, float] = {} last_scan_time: dict[str, float] = {}
# Track market open/close state for notifications
_market_states: dict[str, bool] = {} # market_code -> is_open
shutdown = asyncio.Event() shutdown = asyncio.Event()
def _signal_handler() -> None: def _signal_handler() -> None:
@@ -302,12 +340,31 @@ async def run(settings: Settings) -> None:
logger.info("The Ouroboros is alive. Mode: %s", settings.MODE) logger.info("The Ouroboros is alive. Mode: %s", settings.MODE)
logger.info("Enabled markets: %s", settings.enabled_market_list) logger.info("Enabled markets: %s", settings.enabled_market_list)
# Notify system startup
try:
await telegram.notify_system_start(settings.MODE, settings.enabled_market_list)
except Exception as exc:
logger.warning("System startup notification failed: %s", exc)
try: try:
while not shutdown.is_set(): while not shutdown.is_set():
# Get currently open markets # Get currently open markets
open_markets = get_open_markets(settings.enabled_market_list) open_markets = get_open_markets(settings.enabled_market_list)
if not open_markets: if not open_markets:
# Notify market close for any markets that were open
for market_code, is_open in list(_market_states.items()):
if is_open:
try:
from src.markets.schedule import MARKETS
market_info = MARKETS.get(market_code)
if market_info:
await telegram.notify_market_close(market_info.name, 0.0)
except Exception as exc:
logger.warning("Market close notification failed: %s", exc)
_market_states[market_code] = False
# No markets open — wait until next market opens # No markets open — wait until next market opens
try: try:
next_market, next_open_time = get_next_market_open( next_market, next_open_time = get_next_market_open(
@@ -333,6 +390,14 @@ async def run(settings: Settings) -> None:
if shutdown.is_set(): if shutdown.is_set():
break break
# Notify market open if it just opened
if not _market_states.get(market.code, False):
try:
await telegram.notify_market_open(market.name)
except Exception as exc:
logger.warning("Market open notification failed: %s", exc)
_market_states[market.code] = True
# Volatility Hunter: Scan market periodically to update watchlist # Volatility Hunter: Scan market periodically to update watchlist
now_timestamp = asyncio.get_event_loop().time() now_timestamp = asyncio.get_event_loop().time()
last_scan = last_scan_time.get(market.code, 0.0) last_scan = last_scan_time.get(market.code, 0.0)
@@ -391,12 +456,22 @@ async def run(settings: Settings) -> None:
decision_logger, decision_logger,
context_store, context_store,
criticality_assessor, criticality_assessor,
telegram,
market, market,
stock_code, stock_code,
) )
break # Success — exit retry loop break # Success — exit retry loop
except CircuitBreakerTripped: except CircuitBreakerTripped as exc:
logger.critical("Circuit breaker tripped — shutting down") logger.critical("Circuit breaker tripped — shutting down")
try:
await telegram.notify_circuit_breaker(
pnl_pct=exc.pnl_pct,
threshold=exc.threshold,
)
except Exception as notify_exc:
logger.warning(
"Circuit breaker notification failed: %s", notify_exc
)
raise raise
except ConnectionError as exc: except ConnectionError as exc:
if attempt < MAX_CONNECTION_RETRIES: if attempt < MAX_CONNECTION_RETRIES:

343
tests/test_main.py Normal file
View File

@@ -0,0 +1,343 @@
"""Tests for main trading loop telegram integration."""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.core.risk_manager import CircuitBreakerTripped, FatFingerRejected
from src.main import trading_cycle
class TestTradingCycleTelegramIntegration:
"""Test telegram notifications in trading_cycle function."""
@pytest.fixture
def mock_broker(self) -> MagicMock:
"""Create mock broker."""
broker = MagicMock()
broker.get_orderbook = AsyncMock(
return_value={
"output1": {
"stck_prpr": "50000",
"frgn_ntby_qty": "100",
}
}
)
broker.get_balance = AsyncMock(
return_value={
"output2": [
{
"tot_evlu_amt": "10000000",
"dnca_tot_amt": "5000000",
"pchs_amt_smtl_amt": "5000000",
}
]
}
)
broker.send_order = AsyncMock(return_value={"msg1": "OK"})
return broker
@pytest.fixture
def mock_overseas_broker(self) -> MagicMock:
"""Create mock overseas broker."""
broker = MagicMock()
return broker
@pytest.fixture
def mock_brain(self) -> MagicMock:
"""Create mock brain that decides to buy."""
brain = MagicMock()
decision = MagicMock()
decision.action = "BUY"
decision.confidence = 85
decision.rationale = "Test buy"
brain.decide = AsyncMock(return_value=decision)
return brain
@pytest.fixture
def mock_risk(self) -> MagicMock:
"""Create mock risk manager."""
risk = MagicMock()
risk.validate_order = MagicMock()
return risk
@pytest.fixture
def mock_db(self) -> MagicMock:
"""Create mock database connection."""
return MagicMock()
@pytest.fixture
def mock_decision_logger(self) -> MagicMock:
"""Create mock decision logger."""
logger = MagicMock()
logger.log_decision = MagicMock()
return logger
@pytest.fixture
def mock_context_store(self) -> MagicMock:
"""Create mock context store."""
store = MagicMock()
store.get_latest_timeframe = MagicMock(return_value=None)
return store
@pytest.fixture
def mock_criticality_assessor(self) -> MagicMock:
"""Create mock criticality assessor."""
assessor = MagicMock()
assessor.assess_market_conditions = MagicMock(
return_value=MagicMock(value="NORMAL")
)
assessor.get_timeout = MagicMock(return_value=5.0)
return assessor
@pytest.fixture
def mock_telegram(self) -> MagicMock:
"""Create mock telegram client."""
telegram = MagicMock()
telegram.notify_trade_execution = AsyncMock()
telegram.notify_fat_finger = AsyncMock()
telegram.notify_circuit_breaker = AsyncMock()
return telegram
@pytest.fixture
def mock_market(self) -> MagicMock:
"""Create mock market info."""
market = MagicMock()
market.name = "Korea"
market.code = "KR"
market.exchange_code = "KRX"
market.is_domestic = True
return market
@pytest.mark.asyncio
async def test_trade_execution_notification_sent(
self,
mock_broker: MagicMock,
mock_overseas_broker: MagicMock,
mock_brain: MagicMock,
mock_risk: MagicMock,
mock_db: MagicMock,
mock_decision_logger: MagicMock,
mock_context_store: MagicMock,
mock_criticality_assessor: MagicMock,
mock_telegram: MagicMock,
mock_market: MagicMock,
) -> None:
"""Test telegram notification sent on trade execution."""
with patch("src.main.log_trade"):
await trading_cycle(
broker=mock_broker,
overseas_broker=mock_overseas_broker,
brain=mock_brain,
risk=mock_risk,
db_conn=mock_db,
decision_logger=mock_decision_logger,
context_store=mock_context_store,
criticality_assessor=mock_criticality_assessor,
telegram=mock_telegram,
market=mock_market,
stock_code="005930",
)
# Verify notification was sent
mock_telegram.notify_trade_execution.assert_called_once()
call_kwargs = mock_telegram.notify_trade_execution.call_args.kwargs
assert call_kwargs["stock_code"] == "005930"
assert call_kwargs["market"] == "Korea"
assert call_kwargs["action"] == "BUY"
assert call_kwargs["confidence"] == 85
@pytest.mark.asyncio
async def test_trade_execution_notification_failure_doesnt_crash(
self,
mock_broker: MagicMock,
mock_overseas_broker: MagicMock,
mock_brain: MagicMock,
mock_risk: MagicMock,
mock_db: MagicMock,
mock_decision_logger: MagicMock,
mock_context_store: MagicMock,
mock_criticality_assessor: MagicMock,
mock_telegram: MagicMock,
mock_market: MagicMock,
) -> None:
"""Test trading continues even if notification fails."""
# Make notification fail
mock_telegram.notify_trade_execution.side_effect = Exception("API error")
with patch("src.main.log_trade"):
# Should not raise exception
await trading_cycle(
broker=mock_broker,
overseas_broker=mock_overseas_broker,
brain=mock_brain,
risk=mock_risk,
db_conn=mock_db,
decision_logger=mock_decision_logger,
context_store=mock_context_store,
criticality_assessor=mock_criticality_assessor,
telegram=mock_telegram,
market=mock_market,
stock_code="005930",
)
# Verify notification was attempted
mock_telegram.notify_trade_execution.assert_called_once()
@pytest.mark.asyncio
async def test_fat_finger_notification_sent(
self,
mock_broker: MagicMock,
mock_overseas_broker: MagicMock,
mock_brain: MagicMock,
mock_risk: MagicMock,
mock_db: MagicMock,
mock_decision_logger: MagicMock,
mock_context_store: MagicMock,
mock_criticality_assessor: MagicMock,
mock_telegram: MagicMock,
mock_market: MagicMock,
) -> None:
"""Test telegram notification sent on fat finger rejection."""
# Make risk manager reject the order
mock_risk.validate_order.side_effect = FatFingerRejected(
order_amount=2000000,
total_cash=5000000,
max_pct=30.0,
)
with patch("src.main.log_trade"):
with pytest.raises(FatFingerRejected):
await trading_cycle(
broker=mock_broker,
overseas_broker=mock_overseas_broker,
brain=mock_brain,
risk=mock_risk,
db_conn=mock_db,
decision_logger=mock_decision_logger,
context_store=mock_context_store,
criticality_assessor=mock_criticality_assessor,
telegram=mock_telegram,
market=mock_market,
stock_code="005930",
)
# Verify notification was sent
mock_telegram.notify_fat_finger.assert_called_once()
call_kwargs = mock_telegram.notify_fat_finger.call_args.kwargs
assert call_kwargs["stock_code"] == "005930"
assert call_kwargs["order_amount"] == 2000000
assert call_kwargs["total_cash"] == 5000000
assert call_kwargs["max_pct"] == 30.0
@pytest.mark.asyncio
async def test_fat_finger_notification_failure_still_raises(
self,
mock_broker: MagicMock,
mock_overseas_broker: MagicMock,
mock_brain: MagicMock,
mock_risk: MagicMock,
mock_db: MagicMock,
mock_decision_logger: MagicMock,
mock_context_store: MagicMock,
mock_criticality_assessor: MagicMock,
mock_telegram: MagicMock,
mock_market: MagicMock,
) -> None:
"""Test fat finger exception still raised even if notification fails."""
# Make risk manager reject the order
mock_risk.validate_order.side_effect = FatFingerRejected(
order_amount=2000000,
total_cash=5000000,
max_pct=30.0,
)
# Make notification fail
mock_telegram.notify_fat_finger.side_effect = Exception("API error")
with patch("src.main.log_trade"):
with pytest.raises(FatFingerRejected):
await trading_cycle(
broker=mock_broker,
overseas_broker=mock_overseas_broker,
brain=mock_brain,
risk=mock_risk,
db_conn=mock_db,
decision_logger=mock_decision_logger,
context_store=mock_context_store,
criticality_assessor=mock_criticality_assessor,
telegram=mock_telegram,
market=mock_market,
stock_code="005930",
)
# Verify notification was attempted
mock_telegram.notify_fat_finger.assert_called_once()
@pytest.mark.asyncio
async def test_no_notification_on_hold_decision(
self,
mock_broker: MagicMock,
mock_overseas_broker: MagicMock,
mock_brain: MagicMock,
mock_risk: MagicMock,
mock_db: MagicMock,
mock_decision_logger: MagicMock,
mock_context_store: MagicMock,
mock_criticality_assessor: MagicMock,
mock_telegram: MagicMock,
mock_market: MagicMock,
) -> None:
"""Test no trade notification sent when decision is HOLD."""
# Change brain decision to HOLD
decision = MagicMock()
decision.action = "HOLD"
decision.confidence = 50
decision.rationale = "Insufficient signal"
mock_brain.decide = AsyncMock(return_value=decision)
with patch("src.main.log_trade"):
await trading_cycle(
broker=mock_broker,
overseas_broker=mock_overseas_broker,
brain=mock_brain,
risk=mock_risk,
db_conn=mock_db,
decision_logger=mock_decision_logger,
context_store=mock_context_store,
criticality_assessor=mock_criticality_assessor,
telegram=mock_telegram,
market=mock_market,
stock_code="005930",
)
# Verify no trade notification sent
mock_telegram.notify_trade_execution.assert_not_called()
class TestRunFunctionTelegramIntegration:
"""Test telegram notifications in run function."""
@pytest.mark.asyncio
async def test_circuit_breaker_notification_sent(self) -> None:
"""Test telegram notification sent when circuit breaker trips."""
mock_telegram = MagicMock()
mock_telegram.notify_circuit_breaker = AsyncMock()
# Simulate circuit breaker exception
exc = CircuitBreakerTripped(pnl_pct=-3.5, threshold=-3.0)
# Test the notification logic
try:
await mock_telegram.notify_circuit_breaker(
pnl_pct=exc.pnl_pct,
threshold=exc.threshold,
)
except Exception:
pass # Ignore errors in notification
# Verify notification was called
mock_telegram.notify_circuit_breaker.assert_called_once_with(
pnl_pct=-3.5,
threshold=-3.0,
)