diff --git a/main.py b/main.py index 4d30288..3ed627e 100644 --- a/main.py +++ b/main.py @@ -71,6 +71,7 @@ PlatformRuntimeSettings, _runtime_target_enabled_env, load_platform_runtime_settings, + resolve_telegram_token, ) from strategy_registry import get_platform_profile_status_matrix from strategy_runtime import load_strategy_runtime @@ -127,12 +128,8 @@ def redact_sensitive_text(value: object) -> str: def _get_telegram_token() -> str: - try: - from quant_platform_kit.cloud import get_secret_store - - return get_secret_store().get_secret("firstrade-telegram-token", project_id="firstradequant") - except Exception: - return os.environ.get("TELEGRAM_TOKEN", "") + # Injected unified bot first; legacy firstrade-telegram-token only as last fallback. + return resolve_telegram_token() or "" def _telegram_notification_targets() -> tuple[tuple[str, str], ...]: diff --git a/runtime_config_support.py b/runtime_config_support.py index c32084b..6ec347a 100644 --- a/runtime_config_support.py +++ b/runtime_config_support.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING, Callable, Mapping if TYPE_CHECKING: pass @@ -42,18 +42,50 @@ IBIT_SMART_DCA_PROFILE = "ibit_smart_dca" -def _get_credential(secret_name: str, env_var: str) -> str | None: - val = os.environ.get(env_var) - if val is not None: - return val +FIRSTRADE_SECRET_PROJECT_ID = "firstradequant" +# Legacy per-platform bot; only used when no unified token is injected. +LEGACY_TELEGRAM_TOKEN_SECRET_NAME = "firstrade-telegram-token" + + +def _read_secret(secret_name: str) -> str | None: try: from quant_platform_kit.cloud import get_secret_store - return get_secret_store().get_secret(secret_name, project_id="firstradequant") + return get_secret_store().get_secret(secret_name, project_id=FIRSTRADE_SECRET_PROJECT_ID) except Exception: return None +def _get_credential(secret_name: str, env_var: str) -> str | None: + val = os.environ.get(env_var) + if val is not None: + return val + return _read_secret(secret_name) + + +def resolve_telegram_token( + env: Mapping[str, str] | None = None, + *, + secret_reader: Callable[[str], str | None] | None = None, +) -> str | None: + """Resolve the bot token: injected TELEGRAM_TOKEN, then TELEGRAM_TOKEN_SECRET_NAME, then legacy.""" + + source = os.environ if env is None else env + reader = _read_secret if secret_reader is None else secret_reader + direct = str(source.get("TELEGRAM_TOKEN") or "").strip() + if direct: + return direct + configured_name = str(source.get("TELEGRAM_TOKEN_SECRET_NAME") or "").strip() + names = [configured_name] if configured_name else [] + if LEGACY_TELEGRAM_TOKEN_SECRET_NAME not in names: + names.append(LEGACY_TELEGRAM_TOKEN_SECRET_NAME) + for name in names: + value = str(reader(name) or "").strip() + if value: + return value + return None + + @dataclass(frozen=True) class PlatformRuntimeSettings: project_id: str | None @@ -187,7 +219,7 @@ def load_platform_runtime_settings( strategy_domain=runtime_paths.strategy_domain, strategy_metadata=strategy_metadata, notify_lang=os.getenv("QSL_NOTIFY_LANG") or os.getenv("NOTIFY_LANG", "en"), - tg_token=_get_credential("firstrade-telegram-token", "TELEGRAM_TOKEN"), + tg_token=resolve_telegram_token(), tg_chat_id=os.getenv("QSL_GLOBAL_TELEGRAM_CHAT_ID") or os.getenv("GLOBAL_TELEGRAM_CHAT_ID"), dry_run_only=dry_run_only, runtime_target_enabled=( diff --git a/tests/test_request_handling.py b/tests/test_request_handling.py index 0a5541a..c700ce4 100644 --- a/tests/test_request_handling.py +++ b/tests/test_request_handling.py @@ -17,6 +17,7 @@ def _assume_market_open_for_http_tests(monkeypatch, request): })) monkeypatch.setattr(main, "get_project_id", lambda: "test-project") monkeypatch.setattr("runtime_config_support._get_credential", lambda *_args: None) + monkeypatch.setattr("runtime_config_support._read_secret", lambda *_args: None) if request.node.name == "test_run_endpoint_skips_when_market_closed": return monkeypatch.setattr(main, "_should_skip_for_market_hours", lambda: (False, None)) diff --git a/tests/test_telegram_token_resolution.py b/tests/test_telegram_token_resolution.py new file mode 100644 index 0000000..96280dc --- /dev/null +++ b/tests/test_telegram_token_resolution.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import pytest + +import runtime_config_support +from runtime_config_support import LEGACY_TELEGRAM_TOKEN_SECRET_NAME, resolve_telegram_token + +UNIFIED = "quant-sentinel-telegram-bot-token" + + +def _reader(values): + calls = [] + + def read(name): + calls.append(name) + return values.get(name) + + return read, calls + + +def test_injected_token_wins_over_secret_manager(): + read, calls = _reader({LEGACY_TELEGRAM_TOKEN_SECRET_NAME: "legacy"}) + env = {"TELEGRAM_TOKEN": " unified-token ", "TELEGRAM_TOKEN_SECRET_NAME": UNIFIED} + assert resolve_telegram_token(env, secret_reader=read) == "unified-token" + assert calls == [] + + +def test_configured_secret_name_wins_over_legacy(): + read, calls = _reader({UNIFIED: "unified", LEGACY_TELEGRAM_TOKEN_SECRET_NAME: "legacy"}) + env = {"TELEGRAM_TOKEN": "", "TELEGRAM_TOKEN_SECRET_NAME": UNIFIED} + assert resolve_telegram_token(env, secret_reader=read) == "unified" + assert calls == [UNIFIED] + + +def test_legacy_secret_is_last_fallback(): + read, calls = _reader({LEGACY_TELEGRAM_TOKEN_SECRET_NAME: "legacy"}) + env = {"TELEGRAM_TOKEN_SECRET_NAME": UNIFIED} + assert resolve_telegram_token(env, secret_reader=read) == "legacy" + assert calls == [UNIFIED, LEGACY_TELEGRAM_TOKEN_SECRET_NAME] + + +def test_legacy_only_when_nothing_configured(): + read, calls = _reader({LEGACY_TELEGRAM_TOKEN_SECRET_NAME: "legacy"}) + assert resolve_telegram_token({}, secret_reader=read) == "legacy" + assert calls == [LEGACY_TELEGRAM_TOKEN_SECRET_NAME] + + +def test_missing_everywhere_returns_none(): + read, _ = _reader({}) + assert resolve_telegram_token({"TELEGRAM_TOKEN": " "}, secret_reader=read) is None + + +def test_main_get_telegram_token_prefers_injected(monkeypatch): + pytest.importorskip("flask") + import main + + seen = [] + monkeypatch.setattr(runtime_config_support, "_read_secret", lambda name: seen.append(name) or "legacy") + monkeypatch.setenv("TELEGRAM_TOKEN", "unified-token") + assert main._get_telegram_token() == "unified-token" + assert seen == [] + + +def test_main_get_telegram_token_uses_configured_secret(monkeypatch): + pytest.importorskip("flask") + import main + + values = {UNIFIED: "unified", LEGACY_TELEGRAM_TOKEN_SECRET_NAME: "legacy"} + monkeypatch.setattr(runtime_config_support, "_read_secret", lambda name: values.get(name)) + monkeypatch.delenv("TELEGRAM_TOKEN", raising=False) + monkeypatch.setenv("TELEGRAM_TOKEN_SECRET_NAME", UNIFIED) + assert main._get_telegram_token() == "unified"