fix: reuse configured clients for finance stats

2026-07-03 15:42:29 +08:00
parent 5e99abcbc5
commit a1e2c059b9
4 changed files with 224 additions and 171 deletions
@@ -3,14 +3,14 @@ from __future__ import annotations

 import json
 from datetime import date
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock

 import httpx
-import pytest

 from app.services.finance_service import (
    classify_upstream,
    date_to_shanghai_timestamps,
+    fetch_upstream_cost_sub2api,
    fetch_upstream_cost_new_api,
    fetch_website_revenue,
    get_daily_summary,
@@ -45,6 +45,52 @@ class FakeWebsite:
        self.enabled = True


+class FakeWebsiteClientCtx:
+    def __init__(self, response: dict | None = None, raise_exc: Exception | None = None):
+        self.response = response or {"data": {"total_actual_cost": 123.45}}
+        self.raise_exc = raise_exc
+        self.calls: list[dict] = []
+
+    def __enter__(self): return self
+    def __exit__(self, *a): pass
+
+    def _request(self, method, path, body=None, params=None):
+        self.calls.append({"method": method, "path": path, "body": body, "params": params})
+        if self.raise_exc:
+            raise self.raise_exc
+        return self.response
+
+
+class FakeUpstreamClientCtx:
+    """Context manager mock for UpstreamClient."""
+    def __init__(
+        self,
+        token="tok",
+        user_id="u1",
+        quota_per_unit=500000,
+        response: dict | None = None,
+        raise_exc: Exception | None = None,
+    ):
+        self._token = token
+        self._new_api_user = user_id
+        self._quota_per_unit = quota_per_unit
+        self.response = response or {"data": {"quota": 5_000_000}}
+        self.raise_exc = raise_exc
+        self.calls: list[dict] = []
+
+    def __enter__(self): return self
+    def __exit__(self, *a): pass
+    def ensure_authenticated(self): pass
+    def _new_api_quota_per_unit(self): return self._quota_per_unit
+    def _url(self, path): return f"http://up.test/{path.lstrip('/')}"
+
+    def _send_request(self, method, url, **kwargs):
+        self.calls.append({"method": method, "url": url, **kwargs})
+        if self.raise_exc:
+            raise self.raise_exc
+        return _make_mock_response(self.response)
+
+
 # ─────────────────────────────────────────────
 # classify_upstream tests
 # ─────────────────────────────────────────────
@@ -64,6 +110,15 @@ def test_classify_upstream_nox_token_type():
    assert classify_upstream(u) == "new_api"


+def test_classify_upstream_new_api_by_cookie_mode():
+    u = FakeUpstream(
+        api_prefix="",
+        auth_type="cookie",
+        auth_config_json=json.dumps({"cookie_string": "sid=x", "user_id": "7"}),
+    )
+    assert classify_upstream(u) == "new_api"
+
+
 def test_classify_upstream_new_api_by_login_path():
    u = FakeUpstream(
        api_prefix="",
@@ -120,20 +175,33 @@ def _make_mock_response(json_body: dict, status_code: int = 200):

 def test_website_revenue_success(monkeypatch):
    w = FakeWebsite()
-    mock_resp = _make_mock_response({"data": {"total_actual_cost": 123.45}})
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: mock_resp)
+    client = FakeWebsiteClientCtx({"data": {"total_actual_cost": 123.45}})
+    monkeypatch.setattr("app.services.finance_service.Sub2ApiWebsiteClient", lambda **kw: client)

    amount, err = fetch_website_revenue(w, date(2026, 7, 2))

    assert err is None
    assert abs(amount - 123.45) < 1e-6
+    assert client.calls == [{
+        "method": "GET",
+        "path": "/usage/stats",
+        "body": None,
+        "params": {
+            "start_date": "2026-07-02",
+            "end_date": "2026-07-02",
+            "timezone": "Asia/Shanghai",
+            "nocache": "true",
+        },
+    }]


 def test_website_revenue_nested_data(monkeypatch):
    """total_actual_cost directly in response root (no wrapping data key)."""
    w = FakeWebsite()
-    mock_resp = _make_mock_response({"total_actual_cost": 55.5})
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: mock_resp)
+    monkeypatch.setattr(
+        "app.services.finance_service.Sub2ApiWebsiteClient",
+        lambda **kw: FakeWebsiteClientCtx({"total_actual_cost": 55.5}),
+    )

    amount, err = fetch_website_revenue(w, date(2026, 7, 2))

@@ -143,8 +211,10 @@ def test_website_revenue_nested_data(monkeypatch):

 def test_website_revenue_no_field(monkeypatch):
    w = FakeWebsite()
-    mock_resp = _make_mock_response({"data": {"some_other_field": 99}})
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: mock_resp)
+    monkeypatch.setattr(
+        "app.services.finance_service.Sub2ApiWebsiteClient",
+        lambda **kw: FakeWebsiteClientCtx({"data": {"some_other_field": 99}}),
+    )

    amount, err = fetch_website_revenue(w, date(2026, 7, 2))

@@ -155,11 +225,10 @@ def test_website_revenue_no_field(monkeypatch):

 def test_website_revenue_http_error(monkeypatch):
    w = FakeWebsite()
-    bad_resp = MagicMock()
-    bad_resp.status_code = 403
-    bad_resp.text = "Forbidden"
-    exc = httpx.HTTPStatusError("403 Forbidden", request=MagicMock(), response=bad_resp)
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: (_ for _ in ()).throw(exc))
+    monkeypatch.setattr(
+        "app.services.finance_service.Sub2ApiWebsiteClient",
+        lambda **kw: FakeWebsiteClientCtx(raise_exc=Exception("HTTP 403: Forbidden")),
+    )

    amount, err = fetch_website_revenue(w, date(2026, 7, 2))

@@ -170,7 +239,10 @@ def test_website_revenue_http_error(monkeypatch):

 def test_website_revenue_connection_error(monkeypatch):
    w = FakeWebsite()
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: (_ for _ in ()).throw(Exception("connection refused")))
+    monkeypatch.setattr(
+        "app.services.finance_service.Sub2ApiWebsiteClient",
+        lambda **kw: FakeWebsiteClientCtx(raise_exc=Exception("connection refused")),
+    )

    amount, err = fetch_website_revenue(w, date(2026, 7, 2))

@@ -179,20 +251,25 @@ def test_website_revenue_connection_error(monkeypatch):


 # ─────────────────────────────────────────────
-# fetch_upstream_cost_new_api tests
+# upstream cost tests
 # ─────────────────────────────────────────────

-class FakeUpstreamClientCtx:
-    """Context manager mock for UpstreamClient."""
-    def __init__(self, token="tok", user_id="u1", quota_per_unit=500000):
-        self._token = token
-        self._new_api_user = user_id
-        self._quota_per_unit = quota_per_unit
+def test_upstream_cost_sub2api_uses_upstream_client(monkeypatch):
+    u = FakeUpstream(api_prefix="api/v1", auth_type="bearer")
+    client = FakeUpstreamClientCtx(response={"data": {"total_actual_cost": 42.0}})
+    monkeypatch.setattr("app.services.finance_service.UpstreamClient", lambda **kw: client)

-    def __enter__(self): return self
-    def __exit__(self, *a): pass
-    def ensure_authenticated(self): pass
-    def _new_api_quota_per_unit(self): return self._quota_per_unit
+    amount, err = fetch_upstream_cost_sub2api(u, date(2026, 7, 2))
+
+    assert err is None
+    assert abs(amount - 42.0) < 1e-6
+    assert client.calls[0]["method"] == "GET"
+    assert client.calls[0]["url"] == "http://up.test/usage/stats"
+    assert client.calls[0]["params"] == {
+        "start_date": "2026-07-02",
+        "end_date": "2026-07-02",
+        "timezone": "Asia/Shanghai",
+    }


 def test_upstream_cost_new_api_success(monkeypatch):
@@ -200,21 +277,25 @@ def test_upstream_cost_new_api_success(monkeypatch):
    mock_client_instance = FakeUpstreamClientCtx(quota_per_unit=500000)
    monkeypatch.setattr("app.services.finance_service.UpstreamClient", lambda **kw: mock_client_instance)

-    mock_resp = _make_mock_response({"data": {"quota": 5_000_000}})
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: mock_resp)
-
    amount, err = fetch_upstream_cost_new_api(u, date(2026, 7, 2))

    assert err is None
    assert abs(amount - 10.0) < 1e-6  # 5_000_000 / 500_000 = 10.0
+    assert mock_client_instance.calls[0]["method"] == "GET"
+    assert mock_client_instance.calls[0]["url"] == "http://up.test/api/log/stat"
+    assert mock_client_instance.calls[0]["params"] == {
+        "type": 2,
+        "start_timestamp": 1782921600,
+        "end_timestamp": 1783007999,
+    }


 def test_upstream_cost_new_api_missing_quota(monkeypatch):
    u = FakeUpstream(api_prefix="", auth_type="new_api_token")
-    monkeypatch.setattr("app.services.finance_service.UpstreamClient", lambda **kw: FakeUpstreamClientCtx())
-
-    mock_resp = _make_mock_response({"data": {}})
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: mock_resp)
+    monkeypatch.setattr(
+        "app.services.finance_service.UpstreamClient",
+        lambda **kw: FakeUpstreamClientCtx(response={"data": {}}),
+    )

    amount, err = fetch_upstream_cost_new_api(u, date(2026, 7, 2))

@@ -225,13 +306,15 @@ def test_upstream_cost_new_api_missing_quota(monkeypatch):

 def test_upstream_cost_new_api_http_error(monkeypatch):
    u = FakeUpstream(api_prefix="", auth_type="new_api_token")
-    monkeypatch.setattr("app.services.finance_service.UpstreamClient", lambda **kw: FakeUpstreamClientCtx())

    bad_resp = MagicMock()
    bad_resp.status_code = 401
    bad_resp.text = "Unauthorized"
    exc = httpx.HTTPStatusError("401", request=MagicMock(), response=bad_resp)
-    monkeypatch.setattr("httpx.get", lambda *a, **kw: (_ for _ in ()).throw(exc))
+    monkeypatch.setattr(
+        "app.services.finance_service.UpstreamClient",
+        lambda **kw: FakeUpstreamClientCtx(raise_exc=exc),
+    )

    amount, err = fetch_upstream_cost_new_api(u, date(2026, 7, 2))

@@ -249,16 +332,10 @@ def test_get_daily_summary_website_failure_excluded_from_total(monkeypatch):
    w_ok = FakeWebsite(id=1, name="OK")
    w_fail = FakeWebsite(id=2, name="Fail")

-    call_count = {"n": 0}
-    def mock_get(*a, **kw):
-        call_count["n"] += 1
-        url = a[0] if a else kw.get("url", "")
-        # First call (OK website) returns valid data
-        if call_count["n"] == 1:
-            return _make_mock_response({"data": {"total_actual_cost": 50.0}})
-        # Second call (Fail website) raises
-        raise Exception("timeout")
-    monkeypatch.setattr("httpx.get", mock_get)
+    def mock_fetch_website(row, _date):
+        return (50.0, None) if row.id == 1 else (0.0, "timeout")
+
+    monkeypatch.setattr("app.services.finance_service.fetch_website_revenue", mock_fetch_website)

    result = get_daily_summary([w_ok, w_fail], [], date(2026, 7, 2))

@@ -289,13 +366,8 @@ def test_get_daily_summary_all_success(monkeypatch):
    w = FakeWebsite()
    u = FakeUpstream(api_prefix="api/v1", auth_type="bearer")  # sub2api

-    call_count = {"n": 0}
-    def mock_get(*a, **kw):
-        call_count["n"] += 1
-        if call_count["n"] == 1:
-            return _make_mock_response({"data": {"total_actual_cost": 100.0}})
-        return _make_mock_response({"data": {"total_actual_cost": 40.0}})
-    monkeypatch.setattr("httpx.get", mock_get)
+    monkeypatch.setattr("app.services.finance_service.fetch_website_revenue", lambda *_: (100.0, None))
+    monkeypatch.setattr("app.services.finance_service.fetch_upstream_cost_sub2api", lambda *_: (40.0, None))

    result = get_daily_summary([w], [u], date(2026, 7, 2))