dy/backend/tests/test_worker_scale_controls.py

from __future__ import annotations

import json
import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch


BACKEND_DIR = Path(__file__).resolve().parents[1]
os.environ["KEFU_DB_TYPE"] = "sqlite"
os.environ["KEFU_DATABASE_URL"] = ""
os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db")
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from rpa_engine.douyin_im.session import DouyinImSession
from rpa_engine.playwright_worker import DouyinWorker


class WorkerScaleControlTests(unittest.IsolatedAsyncioTestCase):
    async def test_storage_load_selects_only_cookie_column(self):
        class _ScalarResult:
            def scalar_one_or_none(self):
                return '{"cookies": [{"name": "sessionid", "value": "ok"}]}'

        db = SimpleNamespace(
            execute=AsyncMock(return_value=_ScalarResult()),
            close=AsyncMock(),
        )
        worker = DouyinWorker(account_id=499)
        worker.get_db = AsyncMock(return_value=db)

        storage = await worker._load_storage_state()

        self.assertEqual(storage["cookies"][0]["value"], "ok")
        statement = db.execute.await_args.args[0]
        selected_names = [
            item.get("name") for item in statement.column_descriptions
        ]
        self.assertEqual(selected_names, ["cookie_data"])
        self.assertNotIn("im_session_data", str(statement).lower())

    async def test_direct_service_marks_worker_ready_after_initialization(self):
        worker = DouyinWorker(account_id=500, login_mode="im_direct")
        worker._refresh_follow_welcome_config = AsyncMock(
            return_value=(False, "", "sec-user")
        )
        worker.get_reply_delay = AsyncMock(return_value=None)
        fake_service = SimpleNamespace(run=AsyncMock(), stop=AsyncMock())

        async def complete_initialization():
            service_factory.call_args.kwargs["on_ready"]()

        fake_service.run.side_effect = complete_initialization
        session = DouyinImSession(
            cookies={"sessionid": "session"},
            my_uid=0,
        )

        with patch(
            "rpa_engine.playwright_worker.DouyinImService",
            return_value=fake_service,
        ) as service_factory:
            await worker._run_im_direct_service(session)

        await worker.wait_until_ready()
        worker._refresh_follow_welcome_config.assert_awaited_once_with(force=True)
        fake_service.run.assert_awaited_once_with()
        fake_service.stop.assert_awaited_once_with()

    async def test_direct_start_failure_releases_readiness_waiter(self):
        worker = DouyinWorker(account_id=501, login_mode="im_direct")
        worker._load_storage_state = AsyncMock(return_value=None)
        worker.update_account_status = AsyncMock()
        worker.cleanup = AsyncMock()

        await worker.start()

        with self.assertRaisesRegex(RuntimeError, "未保存 Cookie"):
            await worker.wait_until_ready()

        if worker._task:
            await worker._task
        worker.update_account_status.assert_awaited_once_with(
            "error",
            error_msg="未保存 Cookie，无法直连 IM",
        )

    async def test_prevalidated_start_skips_duplicate_remote_validation(self):
        worker = DouyinWorker(
            account_id=502,
            login_mode="im_direct",
            credential_prevalidated=True,
        )
        session = DouyinImSession(
            cookies={"sessionid": "session"},
            my_uid=10001,
            keys_str=json.dumps({"ec_privateKey": "private"}),
            web_protect_str=json.dumps(
                {
                    "ticket": "ticket",
                    "ts_sign": "sign",
                    "client_cert": "certificate",
                }
            ),
        )
        worker._load_user_agent = AsyncMock(return_value="test-agent")
        worker._build_im_session_from_storage = AsyncMock(return_value=session)
        worker._require_sec_user_id = AsyncMock(return_value="sec-user")
        worker._persist_im_session = AsyncMock()
        worker._run_im_direct_service = AsyncMock()

        with patch(
            "rpa_engine.playwright_worker.validate_im_session",
            new_callable=AsyncMock,
        ) as validate:
            started, reason = await worker._try_cookie_only_im_start(
                {"cookies": []}
            )

        self.assertTrue(started)
        self.assertEqual(reason, "")
        validate.assert_not_awaited()
        worker._require_sec_user_id.assert_awaited_once()
        worker._run_im_direct_service.assert_awaited_once_with(session)

    async def test_disabled_follow_welcome_uses_cached_lightweight_config(self):
        class _Result:
            def first(self):
                return False, "", "sec-user"

        db = SimpleNamespace(
            execute=AsyncMock(return_value=_Result()),
            close=AsyncMock(),
        )
        worker = DouyinWorker(account_id=503)
        worker.get_db = AsyncMock(return_value=db)

        first = await worker._refresh_follow_welcome_config()
        second = await worker._refresh_follow_welcome_config()

        self.assertEqual(first, (False, "", "sec-user"))
        self.assertEqual(second, first)
        worker.get_db.assert_awaited_once_with()
        db.execute.assert_awaited_once()
        db.close.assert_awaited_once_with()

        # The account PUT endpoint calls this synchronous hook so enabling the
        # feature does not wait for the disabled-account ten-minute TTL.
        worker.invalidate_follow_welcome_config()
        await worker._refresh_follow_welcome_config()
        self.assertEqual(worker.get_db.await_count, 2)
        self.assertEqual(db.execute.await_count, 2)

    async def test_disabled_follow_tick_does_not_read_follower_log(self):
        worker = DouyinWorker(account_id=504)
        worker._im_service = SimpleNamespace(session=object())
        worker._follow_config_loaded = True
        worker._refresh_follow_welcome_config = AsyncMock(
            return_value=(False, "", "sec-user")
        )
        worker.get_db = AsyncMock()
        worker._require_sec_user_id = AsyncMock()

        await worker.follow_welcome_tick()

        worker._refresh_follow_welcome_config.assert_awaited_once_with()
        worker.get_db.assert_not_awaited()
        worker._require_sec_user_id.assert_not_awaited()

    async def test_missing_cached_sec_user_id_stops_hosting(self):
        worker = DouyinWorker(account_id=505)
        worker._im_service = SimpleNamespace(session=object())
        worker._follow_config_loaded = True
        worker._refresh_follow_welcome_config = AsyncMock(
            return_value=(False, "", "")
        )
        worker._require_sec_user_id = AsyncMock(return_value="")

        await worker.follow_welcome_tick()

        worker._require_sec_user_id.assert_awaited_once_with(
            "托管运行中"
        )


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