live-hub-py/tests/test_douyu_runner_proxy.py

"""斗鱼任务代理接入三模式行为测试（静态 / API / 关闭）+ 写读分离。"""

import os
import unittest
from unittest.mock import Mock

os.environ.setdefault("DATABASE_URL", "sqlite://")
os.environ.setdefault("APP_ENCRYPTION_KEY", "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=")

from core.douyu.activity_client import DouyuActivityClient
from web.backend.services.douyu_runner import DouyuBatchRunner
from web.backend.services.douyu_runner_core import DouyuBatchRunnerCore, DOUYU_PROXY_TASK_TYPES
from web.backend.models import ProxyConfig


def _cfg(**overrides) -> ProxyConfig:
    values = {
        "enabled": False,
        "api_url": "",
        "http": "",
        "https": "",
        "whitelist_enabled": False,
        "whitelist_platform": "xiequ",
        "whitelist_credentials": None,
        "whitelist_uid": "",
        "whitelist_ukey": "",
    }
    values.update(overrides)
    cfg = ProxyConfig()
    for key, value in values.items():
        setattr(cfg, key, value)
    return cfg


def _runner(db_cfg: ProxyConfig, task_type: str = "query_points"):
    db = Mock()
    db.query.return_value.first.return_value = db_cfg
    runner = DouyuBatchRunner.__new__(DouyuBatchRunner)
    DouyuBatchRunnerCore.__init__(runner, db, "batch", task_type)
    return runner


class DouyuProxyModeTests(unittest.TestCase):
    def test_disabled_mode_goes_direct(self):
        runner = _runner(_cfg(enabled=False), task_type="exchange_goods")
        self.assertIsNone(runner._proxies_for_task())
        self.assertIsNone(runner._proxy_fetcher)
        client = runner._client("acf_uid=100")
        self.assertEqual(client.session.proxies, {})
        # 未开启代理时 execute 客户端不带 proxies 参数
        self.assertIsInstance(client, DouyuActivityClient)

    def test_static_mode_uses_configured_proxies_for_write_tasks(self):
        runner = _runner(_cfg(enabled=True, http="http://10.0.0.1:8080"), task_type="exchange_goods")
        proxies = runner._proxies_for_task()
        self.assertEqual(proxies, {"http": "http://10.0.0.1:8080", "https": "http://10.0.0.1:8080"})
        client = runner._client("acf_uid=100")
        self.assertEqual(client.session.proxies["https"], "http://10.0.0.1:8080")
        # 每次任务都取同一静态代理
        self.assertEqual(runner._proxies_for_task(), proxies)

    def test_api_mode_fetches_new_proxy_per_write_task(self):
        cfg = _cfg(enabled=True, api_url="https://proxy-api.example/fetch")
        runner = _runner(cfg, task_type="exchange_goods")
        self.assertIsNotNone(runner._proxy_fetcher)
        fake = Mock()
        fake.fetch_new_proxy.return_value = "http://10.0.0.2:3128"
        runner._proxy_fetcher = fake

        proxies = runner._proxies_for_task()
        self.assertEqual(proxies, {"http": "http://10.0.0.2:3128", "https": "http://10.0.0.2:3128"})
        client = runner._client("acf_uid=100")
        self.assertEqual(client.session.proxies["http"], "http://10.0.0.2:3128")

        # 每任务取新代理：_client 也会再 fetch 一次（客户端构造即取代理）
        fake.fetch_new_proxy.return_value = "http://10.0.0.3:3128"
        self.assertEqual(runner._proxies_for_task()["http"], "http://10.0.0.3:3128")
        self.assertEqual(fake.fetch_new_proxy.call_count, 3)

    def test_api_mode_failure_falls_back_to_direct(self):
        runner = _runner(_cfg(enabled=True, api_url="https://proxy-api.example/fetch"), task_type="exchange_goods")
        runner._proxy_fetcher = Mock()
        runner._proxy_fetcher.fetch_new_proxy.return_value = None
        runner._push_log = Mock()

        self.assertIsNone(runner._proxies_for_task())
        client = runner._client("acf_uid=100")
        self.assertEqual(client.session.proxies, {})
        runner._push_log.assert_called()

    def test_read_tasks_stay_direct_even_with_proxy_enabled(self):
        """写读分离：查询/刷新/绑定等读任务即使代理开启也直连，不消耗代理额度。"""
        runner = _runner(_cfg(enabled=True, http="http://10.0.0.1:8080"), task_type="query_points")
        self.assertIsNotNone(runner._static_proxies)
        self.assertIsNone(runner._proxies_for_task())
        client = runner._client("acf_uid=100")
        self.assertEqual(client.session.proxies, {})

        # API 模式同理：读任务不触发 fetch
        runner = _runner(_cfg(enabled=True, api_url="https://proxy-api.example/fetch"), task_type="query_exchange_records")
        runner._proxy_fetcher = Mock()
        runner._proxy_fetcher.fetch_new_proxy.side_effect = AssertionError("读任务不应取代理")
        self.assertIsNone(runner._proxies_for_task())

    def test_proxy_task_types_covers_all_write_tasks(self):
        """白名单集合应覆盖全部会消耗/下单的写操作任务，且不含查询类。"""
        self.assertIn("exchange_goods", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("lock_goods", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("pay_locked_order", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("exchange_esports_goods", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("exchange_xpd_goods", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("create_gold_qr", DOUYU_PROXY_TASK_TYPES)
        self.assertIn("donate_elite_gift", DOUYU_PROXY_TASK_TYPES)
        self.assertNotIn("query_points", DOUYU_PROXY_TASK_TYPES)
        self.assertNotIn("refresh_goods", DOUYU_PROXY_TASK_TYPES)
        self.assertNotIn("get_bind_qr", DOUYU_PROXY_TASK_TYPES)

    def test_activity_client_accepts_proxies_at_session_level(self):
        client = DouyuActivityClient(
            "acf_uid=100",
            proxies={"http": "http://10.0.0.9:7890", "https": "http://10.0.0.9:7890"},
        )
        self.assertEqual(client.session.proxies["http"], "http://10.0.0.9:7890")
        # trust_env=False 维持不变：不受环境变量代理影响
        self.assertFalse(client.session.trust_env)


if __name__ == "__main__":
    unittest.main()