exec-mcp/tests/test_reporting_tools.py

from __future__ import annotations

import json

import pytest

from exec_mcp import repo, storage
from exec_mcp import services_bitstamp, services_orders
from exec_mcp.server import report_balance_snapshots, report_recent_trades


@pytest.fixture()
def temp_db(tmp_path, monkeypatch):
    db_path = tmp_path / "exec.sqlite3"
    monkeypatch.setattr(storage, "DB_PATH", db_path)
    storage.init_db()
    return db_path


def _create_account() -> str:
    repo.create_account(display_name="strategy", venue="bitstamp", venue_account_ref="ref-1", api_key="k", api_secret="s")
    return repo.list_accounts(enabled_only=False)[0]["id"]


def test_report_recent_trades_defaults_to_finished_and_sorts_newest_first(temp_db):
    account_id = _create_account()

    with storage.get_connection() as conn:
        rows = [
            ("rec-1", "finished", "client-a", "1994292738961401", "2026-04-09T08:20:50+00:00"),
            ("rec-2", "cancelled", "client-a", "1994292738961402", "2026-04-09T09:20:50+00:00"),
            ("rec-3", "finished", "client-b", "1994292738961403", "2026-04-09T10:20:50+00:00"),
            ("rec-4", "open", "client-a", "1994292738961404", "2026-04-09T11:20:50+00:00"),
            ("rec-5", "new", "client-a", "1994292738961405", "2026-04-09T12:20:50+00:00"),
        ]
        for record_id, status, client_id, order_id, stamp in rows:
            conn.execute(
                """
                INSERT INTO order_records
                (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    record_id,
                    account_id,
                    "xrpusd",
                    "buy",
                    "limit",
                    "10",
                    "2.00",
                    None,
                    status,
                    order_id,
                    client_id,
                    None,
                    json.dumps({"id": order_id, "status": status}),
                    stamp,
                    stamp,
                ),
            )
        conn.commit()

    result = report_recent_trades(account_id=account_id)
    assert result["ok"] is True
    assert result["status"] == "finished"
    assert result["count"] == 2
    assert [row["bitstamp_order_id"] for row in result["orders"]] == ["1994292738961403", "1994292738961401"]


def test_report_recent_trades_supports_open_status_and_client_filter(temp_db):
    account_id = _create_account()

    with storage.get_connection() as conn:
        for idx, (status, client_id, stamp) in enumerate(
            [
                ("open", "client-a", "2026-04-09T08:20:50+00:00"),
                ("new", "client-a", "2026-04-09T09:20:50+00:00"),
                ("partially_filled", "client-b", "2026-04-09T10:20:50+00:00"),
                ("cancelled", "client-a", "2026-04-09T11:20:50+00:00"),
            ],
            start=1,
        ):
            order_id = f"19942927389614{idx}"
            conn.execute(
                """
                INSERT INTO order_records
                (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    f"rec-{idx}",
                    account_id,
                    "xrpusd",
                    "sell",
                    "limit",
                    "10",
                    "2.00",
                    None,
                    status,
                    order_id,
                    client_id,
                    None,
                    json.dumps({"id": order_id, "status": status}),
                    stamp,
                    stamp,
                ),
            )
        conn.commit()

    result = report_recent_trades(account_id=account_id, status="open", client_id="client-a", limit=10)
    assert result["ok"] is True
    assert result["status"] == "open"
    assert result["client_id"] == "client-a"
    assert [row["status"] for row in result["orders"]] == ["new", "open"]


def test_report_recent_trades_supports_all_statuses_and_client_filter(temp_db):
    account_id = _create_account()

    with storage.get_connection() as conn:
        for idx, (status, client_id, stamp) in enumerate(
            [
                ("finished", "client-a", "2026-04-09T08:20:50+00:00"),
                ("cancelled", "client-a", "2026-04-09T09:20:50+00:00"),
                ("finished", "client-b", "2026-04-09T10:20:50+00:00"),
            ],
            start=1,
        ):
            order_id = f"19942927389614{idx}"
            conn.execute(
                """
                INSERT INTO order_records
                (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    f"rec-all-{idx}",
                    account_id,
                    "xrpusd",
                    "sell",
                    "limit",
                    "10",
                    "2.00",
                    None,
                    status,
                    order_id,
                    client_id,
                    None,
                    json.dumps({"id": order_id, "status": status}),
                    stamp,
                    stamp,
                ),
            )
        conn.commit()

    result = report_recent_trades(account_id=account_id, status="all", client_id="client-a", limit=10)
    assert result["ok"] is True
    assert result["status"] == "all"
    assert result["client_id"] == "client-a"
    assert [row["status"] for row in result["orders"]] == ["cancelled", "finished"]


def test_report_balance_snapshots_returns_history_newest_first(temp_db):
    account_id = _create_account()

    repo.save_account_balance_snapshot(
        account_id=account_id,
        source_kind="order_finished",
        source_ref="order-1",
        snapshot={"balances": [{"asset_code": "BTC"}], "total_value_usd": 11250.0},
    )
    repo.save_account_balance_snapshot(
        account_id=account_id,
        source_kind="order_finished",
        source_ref="order-2",
        snapshot={"balances": [{"asset_code": "USD"}], "total_value_usd": 500.0},
    )

    result = report_balance_snapshots(account_id=account_id)
    assert len(result) == 2
    assert [row["source_ref"] for row in result] == ["order-2", "order-1"]
    assert [row["snapshot"]["balances"][0]["asset_code"] for row in result] == ["USD", "BTC"]


def test_record_order_status_update_persists_raw_json_and_triggers_snapshot_once(temp_db, monkeypatch):
    account_id = _create_account()

    with storage.get_connection() as conn:
        conn.execute(
            """
            INSERT INTO order_records
            (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                "rec-1",
                account_id,
                "xrpusd",
                "buy",
                "limit",
                "10",
                "2.00",
                None,
                "open",
                "1994292738961401",
                "client-a",
                None,
                json.dumps({"status": "open"}),
                "2026-04-09T08:20:50+00:00",
                "2026-04-09T08:20:50+00:00",
            ),
        )
        conn.commit()

    snapshot_calls: list[dict] = []
    monkeypatch.setattr(
        services_orders,
        "capture_account_balance_snapshot",
        lambda account_id, *, source_kind, source_ref="": snapshot_calls.append(
            {"account_id": account_id, "source_kind": source_kind, "source_ref": source_ref}
        ),
    )

    services_orders._record_order_status_update(
        bitstamp_order_id="1994292738961401",
        status="finished",
        raw_json={"id": "1994292738961401", "status": "finished"},
        account_id=account_id,
    )

    with storage.get_connection() as conn:
        row = conn.execute(
            "SELECT status, raw_json FROM order_records WHERE bitstamp_order_id = ?",
            ("1994292738961401",),
        ).fetchone()

    assert row["status"] == "finished"
    assert json.loads(row["raw_json"]) == {"id": "1994292738961401", "status": "finished"}
    assert snapshot_calls == [
        {"account_id": account_id, "source_kind": "order_finished", "source_ref": "1994292738961401"}
    ]

    services_orders._record_order_status_update(
        bitstamp_order_id="1994292738961401",
        status="finished",
        raw_json={"id": "1994292738961401", "status": "finished"},
        account_id=account_id,
    )
    assert len(snapshot_calls) == 1


def test_record_order_status_update_ws_path_defers_snapshot(temp_db, monkeypatch):
    """When account_id is not passed (WS-driven call), the balance snapshot
    is deferred to the background pool instead of running synchronously."""
    account_id = _create_account()

    with storage.get_connection() as conn:
        conn.execute(
            """
            INSERT INTO order_records
            (id, account_id, market, side, order_type, amount, price, expire_time, status, bitstamp_order_id, client_id, client_order_id, raw_json, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                "rec-ws-1",
                account_id,
                "xrpusd",
                "buy",
                "limit",
                "10",
                "2.00",
                None,
                "open",
                "1994292738961402",
                "client-a",
                None,
                json.dumps({"status": "open"}),
                "2026-04-09T08:20:50+00:00",
                "2026-04-09T08:20:50+00:00",
            ),
        )
        conn.commit()

    submitted: list[tuple] = []

    class _FakePool:
        def submit(self, fn, *a):
            submitted.append((fn, a))

    monkeypatch.setattr(services_orders, "_SNAPSHOT_POOL", _FakePool())

    # No account_id — simulates the WS handler call path
    result = services_orders._record_order_status_update(
        bitstamp_order_id="1994292738961402",
        status="finished",
        raw_json={"id": "1994292738961402", "status": "finished"},
    )

    assert result["ok"] is True
    assert result["status"] == "finished"
    assert len(submitted) == 1
    assert submitted[0][0] is services_orders._capture_snapshot_task
    assert submitted[0][1] == (account_id, "1994292738961402")


def test_capture_account_balance_snapshot_uses_crypto_price(temp_db, monkeypatch):
    account_id = _create_account()

    monkeypatch.setattr(
        services_bitstamp,
        "fetch_account_balance",
        lambda _account_id: {
            "source": "bitstamp",
            "cached": False,
            "balances": [
                {
                    "account_id": _account_id,
                    "asset_code": "XRP",
                    "available": 1.0,
                    "reserved": 0.0,
                    "total": 2.0,
                }
            ],
            "payload": [{"currency": "xrp", "total": "2.0", "available": "1.0", "reserved": "0.0"}],
        },
    )
    monkeypatch.setattr(
        services_bitstamp,
        "get_crypto_price",
        lambda symbol: {"symbol": symbol.upper(), "price": 1.4109, "timestamp": 1779039021},
    )

    result = services_bitstamp.capture_account_balance_snapshot(
        account_id,
        source_kind="order_finished",
        source_ref="order-crypto-price",
    )

    assert result["ok"] is True
    snapshot = result["snapshot"]
    assert snapshot["balances"][0]["price_usd"] == 1.4109
    with storage.get_connection() as conn:
        row = conn.execute(
            "SELECT snapshot_json FROM account_balance_snapshots WHERE account_id = ?",
            (account_id,),
        ).fetchone()

    stored = json.loads(row["snapshot_json"])
    assert stored["balances"][0]["price_usd"] == 1.4109


def test_get_account_info_stays_available_when_crypto_price_is_unavailable(temp_db, monkeypatch):
    account_id = _create_account()

    monkeypatch.setattr(
        services_bitstamp,
        "fetch_account_balance",
        lambda _account_id: {
            "source": "bitstamp",
            "cached": False,
            "balances": [
                {
                    "account_id": _account_id,
                    "asset_code": "XRP",
                    "available": 1.5,
                    "reserved": 0.5,
                    "total": 2.0,
                },
                {
                    "account_id": _account_id,
                    "asset_code": "USD",
                    "available": 10.0,
                    "reserved": 5.0,
                    "total": 15.0,
                },
            ],
            "payload": [],
        },
    )
    monkeypatch.setattr(services_bitstamp, "get_crypto_price", lambda symbol: {"error": "TIMEOUT", "symbol": symbol})

    result = services_bitstamp.fetch_account_info(account_id)

    assert result["id"] == account_id
    assert result["total_value_usd"] == 15.0
    assert result["balances"][0]["asset_code"] == "XRP"
    assert result["balances"][0]["price_usd"] is None
    assert result["balances"][0]["value_usd"] is None
    assert result["balances"][1]["asset_code"] == "USD"
    assert result["balances"][1]["price_usd"] == 1.0
    assert result["balances"][1]["value_usd"] == 15.0


def test_get_account_info_falls_back_when_balance_fetch_fails(temp_db, monkeypatch):
    account_id = _create_account()

    monkeypatch.setattr(services_bitstamp, "fetch_account_balance", lambda _account_id: (_ for _ in ()).throw(RuntimeError("bitstamp unavailable")))
    monkeypatch.setattr(services_bitstamp, "get_crypto_price", lambda symbol: {"symbol": symbol.upper(), "price": 1.4109, "timestamp": 1779039021})

    result = services_bitstamp.fetch_account_info(account_id)

    assert result["id"] == account_id
    assert result["balances"] == []
    assert result["total_value_usd"] == 0.0
    assert result["raw_balance"] is None
    assert "balance_error" in result