diff --git a/src/main.py b/src/main.py index ba14d64..e32701d 100644 --- a/src/main.py +++ b/src/main.py @@ -606,8 +606,37 @@ async def run(settings: Settings) -> None: ) await telegram.send_message(message) + async def handle_stop() -> None: + """Handle /stop command - pause trading.""" + if not pause_trading.is_set(): + await telegram.send_message("⏸️ Trading is already paused") + return + + pause_trading.clear() + logger.info("Trading paused via Telegram command") + await telegram.send_message( + "⏸️ Trading Paused\n\n" + "All trading operations have been suspended.\n" + "Use /resume to restart trading." + ) + + async def handle_resume() -> None: + """Handle /resume command - resume trading.""" + if pause_trading.is_set(): + await telegram.send_message("▶️ Trading is already active") + return + + pause_trading.set() + logger.info("Trading resumed via Telegram command") + await telegram.send_message( + "▶️ Trading Resumed\n\n" + "Trading operations have been restarted." + ) + command_handler.register_command("start", handle_start) command_handler.register_command("help", handle_help) + command_handler.register_command("stop", handle_stop) + command_handler.register_command("resume", handle_resume) # Initialize volatility hunter volatility_analyzer = VolatilityAnalyzer(min_volume_surge=2.0, min_price_change=1.0) @@ -636,7 +665,10 @@ async def run(settings: Settings) -> None: # Track market open/close state for notifications _market_states: dict[str, bool] = {} # market_code -> is_open + # Trading control events shutdown = asyncio.Event() + pause_trading = asyncio.Event() + pause_trading.set() # Default: trading enabled def _signal_handler() -> None: logger.info("Shutdown signal received") @@ -674,6 +706,9 @@ async def run(settings: Settings) -> None: session_interval = settings.SESSION_INTERVAL_HOURS * 3600 # Convert to seconds while not shutdown.is_set(): + # Wait for trading to be unpaused + await pause_trading.wait() + try: await run_daily_session( broker, @@ -706,6 +741,9 @@ async def run(settings: Settings) -> None: logger.info("Realtime trading mode: 60s interval per stock") while not shutdown.is_set(): + # Wait for trading to be unpaused + await pause_trading.wait() + # Get currently open markets open_markets = get_open_markets(settings.enabled_market_list) diff --git a/tests/test_telegram_commands.py b/tests/test_telegram_commands.py index 91332dd..978ef45 100644 --- a/tests/test_telegram_commands.py +++ b/tests/test_telegram_commands.py @@ -253,6 +253,195 @@ class TestUpdateHandling: await handler._handle_update(update) +class TestTradingControlCommands: + """Test trading control commands.""" + + @pytest.mark.asyncio + async def test_stop_command_pauses_trading(self) -> None: + """Stop command clears pause event.""" + client = TelegramClient(bot_token="123:abc", chat_id="456", enabled=True) + handler = TelegramCommandHandler(client) + + # Create mock pause event + import asyncio + + pause_event = asyncio.Event() + pause_event.set() # Initially active + + mock_resp = AsyncMock() + mock_resp.status = 200 + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + async def mock_stop() -> None: + """Mock /stop handler.""" + if not pause_event.is_set(): + await client.send_message("⏸️ Trading is already paused") + return + + pause_event.clear() + await client.send_message( + "⏸️ Trading Paused\n\n" + "All trading operations have been suspended.\n" + "Use /resume to restart trading." + ) + + handler.register_command("stop", mock_stop) + + with patch("aiohttp.ClientSession.post", return_value=mock_resp) as mock_post: + update = { + "update_id": 1, + "message": { + "chat": {"id": 456}, + "text": "/stop", + }, + } + + await handler._handle_update(update) + + # Verify pause event was cleared + assert not pause_event.is_set() + + # Verify message was sent + assert mock_post.call_count == 1 + payload = mock_post.call_args.kwargs["json"] + assert "Trading Paused" in payload["text"] + + @pytest.mark.asyncio + async def test_resume_command_resumes_trading(self) -> None: + """Resume command sets pause event.""" + client = TelegramClient(bot_token="123:abc", chat_id="456", enabled=True) + handler = TelegramCommandHandler(client) + + # Create mock pause event (initially paused) + import asyncio + + pause_event = asyncio.Event() + pause_event.clear() # Initially paused + + mock_resp = AsyncMock() + mock_resp.status = 200 + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + async def mock_resume() -> None: + """Mock /resume handler.""" + if pause_event.is_set(): + await client.send_message("▶️ Trading is already active") + return + + pause_event.set() + await client.send_message( + "▶️ Trading Resumed\n\n" + "Trading operations have been restarted." + ) + + handler.register_command("resume", mock_resume) + + with patch("aiohttp.ClientSession.post", return_value=mock_resp) as mock_post: + update = { + "update_id": 1, + "message": { + "chat": {"id": 456}, + "text": "/resume", + }, + } + + await handler._handle_update(update) + + # Verify pause event was set + assert pause_event.is_set() + + # Verify message was sent + assert mock_post.call_count == 1 + payload = mock_post.call_args.kwargs["json"] + assert "Trading Resumed" in payload["text"] + + @pytest.mark.asyncio + async def test_stop_when_already_paused(self) -> None: + """Stop command when already paused sends appropriate message.""" + client = TelegramClient(bot_token="123:abc", chat_id="456", enabled=True) + handler = TelegramCommandHandler(client) + + # Create mock pause event (already paused) + import asyncio + + pause_event = asyncio.Event() + pause_event.clear() + + mock_resp = AsyncMock() + mock_resp.status = 200 + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + async def mock_stop() -> None: + """Mock /stop handler.""" + if not pause_event.is_set(): + await client.send_message("⏸️ Trading is already paused") + return + + pause_event.clear() + + handler.register_command("stop", mock_stop) + + with patch("aiohttp.ClientSession.post", return_value=mock_resp) as mock_post: + update = { + "update_id": 1, + "message": { + "chat": {"id": 456}, + "text": "/stop", + }, + } + + await handler._handle_update(update) + + # Verify message was sent + payload = mock_post.call_args.kwargs["json"] + assert "already paused" in payload["text"] + + @pytest.mark.asyncio + async def test_resume_when_already_active(self) -> None: + """Resume command when already active sends appropriate message.""" + client = TelegramClient(bot_token="123:abc", chat_id="456", enabled=True) + handler = TelegramCommandHandler(client) + + # Create mock pause event (already active) + import asyncio + + pause_event = asyncio.Event() + pause_event.set() + + mock_resp = AsyncMock() + mock_resp.status = 200 + mock_resp.__aenter__ = AsyncMock(return_value=mock_resp) + mock_resp.__aexit__ = AsyncMock(return_value=False) + + async def mock_resume() -> None: + """Mock /resume handler.""" + if pause_event.is_set(): + await client.send_message("▶️ Trading is already active") + return + + pause_event.set() + + handler.register_command("resume", mock_resume) + + with patch("aiohttp.ClientSession.post", return_value=mock_resp) as mock_post: + update = { + "update_id": 1, + "message": { + "chat": {"id": 456}, + "text": "/resume", + }, + } + + await handler._handle_update(update) + + # Verify message was sent + payload = mock_post.call_args.kwargs["json"] + assert "already active" in payload["text"] + + class TestBasicCommands: """Test basic command implementations."""