diff --git a/src/main.py b/src/main.py index 324ef54..b303ad8 100644 --- a/src/main.py +++ b/src/main.py @@ -10,6 +10,7 @@ import argparse import asyncio import logging import signal +import sys from datetime import UTC, datetime from typing import Any @@ -23,11 +24,12 @@ from src.context.layer import ContextLayer from src.context.store import ContextStore from src.core.criticality import CriticalityAssessor 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.logging.decision_logger import DecisionLogger from src.logging_config import setup_logging from src.markets.schedule import MarketInfo, get_next_market_open, get_open_markets +from src.notifications.telegram_client import TelegramClient logger = logging.getLogger(__name__) @@ -62,6 +64,7 @@ async def trading_cycle( decision_logger: DecisionLogger, context_store: ContextStore, criticality_assessor: CriticalityAssessor, + telegram: TelegramClient, market: MarketInfo, stock_code: str, ) -> None: @@ -199,11 +202,23 @@ async def trading_cycle( order_amount = current_price * quantity # 4. Risk check BEFORE order - risk.validate_order( - current_pnl_pct=pnl_pct, - order_amount=order_amount, - total_cash=total_cash, - ) + try: + risk.validate_order( + current_pnl_pct=pnl_pct, + order_amount=order_amount, + 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 if market.is_domestic: @@ -223,6 +238,19 @@ async def trading_cycle( ) 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 log_trade( conn=db_conn, @@ -266,6 +294,13 @@ async def run(settings: Settings) -> None: decision_logger = DecisionLogger(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 volatility_analyzer = VolatilityAnalyzer(min_volume_surge=2.0, min_price_change=1.0) market_scanner = MarketScanner( @@ -289,6 +324,9 @@ async def run(settings: Settings) -> None: # Track last scan time for each market 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() 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("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: while not shutdown.is_set(): # Get currently open markets open_markets = get_open_markets(settings.enabled_market_list) 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 try: next_market, next_open_time = get_next_market_open( @@ -333,6 +390,14 @@ async def run(settings: Settings) -> None: if shutdown.is_set(): 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 now_timestamp = asyncio.get_event_loop().time() last_scan = last_scan_time.get(market.code, 0.0) @@ -391,12 +456,22 @@ async def run(settings: Settings) -> None: decision_logger, context_store, criticality_assessor, + telegram, market, stock_code, ) break # Success — exit retry loop - except CircuitBreakerTripped: + except CircuitBreakerTripped as exc: 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 except ConnectionError as exc: if attempt < MAX_CONNECTION_RETRIES: diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..62139da --- /dev/null +++ b/tests/test_main.py @@ -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, + )