alphax/tests/test_chat_assistant.py

from datetime import datetime, timedelta

import pandas as pd
from fastapi.testclient import TestClient

from app.db import auth_db, chat_assistant_db
from app.services import chat_assistant
from app.web import web_server


def _login_user(email: str = "chat-user@example.com", password: str = "StrongPass123", monkeypatch=None) -> str:
    if monkeypatch is not None:
        monkeypatch.setattr(auth_db, "is_smtp_configured", lambda: False)
    reg = auth_db.register_user(email, password)
    auth_db.verify_email(email, reg["verification_code"])
    user = auth_db.get_user_by_email(email)
    auth_db.claim_free_trial(user["id"])
    return auth_db.login_user(email, password)["token"]


def _fake_ohlcv(rows=180, start=100.0):
    now = datetime(2026, 5, 1, 12, 0, 0)
    data = []
    price = start
    for i in range(rows):
        ts = now + timedelta(minutes=i * 15)
        open_p = price
        close_p = price * (1 + 0.001)
        high = max(open_p, close_p) * 1.004
        low = min(open_p, close_p) * 0.996
        vol = 1000 + i * 3
        data.append([ts, open_p, high, low, close_p, vol])
        price = close_p
    return pd.DataFrame(data, columns=["timestamp", "open", "high", "low", "close", "volume"])


def test_chat_page_and_bootstrap_require_subscription():
    token = _login_user("chat-page@example.com")
    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)

    page = client.get("/chat")
    boot = client.get("/api/chat/bootstrap")

    assert page.status_code == 200
    assert "Crypto 研究助手" in page.text
    assert boot.status_code == 200
    assert "suggested_prompts" in boot.json()


def test_single_coin_chat_fetches_multi_timeframe_technicals(monkeypatch):
    token = _login_user("chat-coin@example.com", monkeypatch=monkeypatch)
    seen = []

    def fake_fetch(symbol, timeframe, limit=160):
        seen.append((symbol, timeframe))
        return _fake_ohlcv(rows=180, start=100.0)

    monkeypatch.setattr(chat_assistant, "fetch_binance_klines", fake_fetch)
    monkeypatch.setattr(chat_assistant, "_call_chat_llm", lambda message, context, history=None: {"status": "skipped", "error": "disabled"})

    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)
    resp = client.post("/api/chat/send", json={"message": "分析 SUI/USDT 现在的技术面"})

    assert resp.status_code == 200
    data = resp.json()
    assert data["intent"] == "coin_analysis"
    assert data["assistant_message"]["content"]["answer_style"] == "technical"
    assert data["symbol"] == "SUI/USDT"
    assert {tf for _, tf in seen} >= {"15m", "1h", "4h", "1d"}
    tfs = data["assistant_message"]["context"]["technicals"]["timeframes"]
    assert tfs["15m"]["available"] is True
    assert tfs["1h"]["pa"]["zone_count"] >= 0


def test_non_crypto_question_is_rejected(monkeypatch):
    token = _login_user("chat-scope@example.com", monkeypatch=monkeypatch)
    monkeypatch.setattr(chat_assistant, "_call_chat_llm", lambda message, context, history=None: {"status": "skipped", "error": "disabled"})
    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)

    resp = client.post("/api/chat/send", json={"message": "帮我写一份旅游攻略"})

    assert resp.status_code == 200
    data = resp.json()
    assert data["intent"] == "unsupported"
    assert "只能回答加密货币" in data["answer"]["summary"]


def test_chat_rejects_paper_trading_questions(monkeypatch):
    token = _login_user("chat-restricted@example.com", monkeypatch=monkeypatch)
    monkeypatch.setattr(chat_assistant, "_call_chat_llm", lambda message, context, history=None: {"status": "skipped", "error": "disabled"})
    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)

    resp = client.post("/api/chat/send", json={"message": "帮我看一下策略交易开仓和收益"})

    assert resp.status_code == 200
    data = resp.json()
    assert data["intent"] == "restricted"
    assert "内部策略交易数据不可在智能问答中直接访问" in data["answer"]["summary"]
    assert data["answer"]["evidence"] == []


def test_chat_user_memory_tracks_last_symbol(monkeypatch):
    token = _login_user("chat-memory@example.com", monkeypatch=monkeypatch)
    user = auth_db.get_user_by_email("chat-memory@example.com")
    monkeypatch.setattr(chat_assistant, "fetch_binance_klines", lambda symbol, timeframe, limit=160: _fake_ohlcv())
    monkeypatch.setattr(chat_assistant, "_call_chat_llm", lambda message, context, history=None: {"status": "skipped", "error": "disabled"})

    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)
    client.post("/api/chat/send", json={"message": "分析 LINK/USDT 技术面"})
    resp = client.post("/api/chat/send", json={"message": "那它现在追高风险大吗？"})

    prefs = chat_assistant_db.get_user_preferences(user["id"])
    assert prefs["last_symbol"] == "LINK/USDT"
    assert resp.json()["symbol"] == "LINK/USDT"


def test_chat_repairs_llm_mojibake_output(monkeypatch):
    token = _login_user("chat-mojibake@example.com", monkeypatch=monkeypatch)
    monkeypatch.setattr(chat_assistant, "fetch_binance_klines", lambda symbol, timeframe, limit=160: _fake_ohlcv())
    monkeypatch.setattr(
        chat_assistant,
        "_call_chat_llm",
        lambda message, context, history=None: {
            "status": "success",
            "model": "mock",
            "content": {
                "summary": "BTC/USDT 当前技术面偏弱",
                "answer": "结论：BTC 短线偏弱",
                "evidence": ["证据：1h 低于 MA20"],
            },
        },
    )

    client = TestClient(web_server.app)
    client.cookies.set("altcoin_session", token)
    resp = client.post("/api/chat/send", json={"message": "分析 BTC/USDT 技术面"})

    assert resp.status_code == 200
    data = resp.json()
    content = data["assistant_message"]["content"]
    assert "当前技术面偏弱" in content["summary"]
    assert "结论：BTC 短线偏弱" in content["answer"]
    assert "证据：1h 低于 MA20" in content["evidence"][0]