creator-hub/browser_gateway/test_private_messages.py

import unittest
import threading
import json
import subprocess
from contextlib import nullcontext, contextmanager
from unittest.mock import Mock, patch

from .platform.douyin import DouyinBrowser, DouyinError, inbox_expression, im_expression, private_identity_expression
from .server.http import Gateway, GatewayHandler, RequestError


class PrivateMessageTests(unittest.TestCase):
    def test_private_login_uses_sdk_uid_and_never_queries_profile(self):
        script='''
const vm=require('node:vm'),assert=require('node:assert/strict');
const source=SOURCE;
const sdk={ctx:{option:{userId:'123'}},getConversationList:()=>{throw Error('login must not read messages')}};
const req={m:{one:function(){/* getOrCreatePrivateConversationByUid */}},c:{one:{exports:{A:{instance:{imSdkService:{imSdkManager:{getImSdkInstance:()=>sdk}}}}}}}};
const chunks=[];chunks.push=function(args){args[2](req)};
const context={location:{origin:'https://www.douyin.com'},fetch:()=>{throw Error('profile is not chat authentication')},crypto:{randomUUID:()=> 'id'},window:{'@pc-im/im:version':chunks}};
(async()=>{
 let result=await vm.runInNewContext(source,context);assert.equal(result.status,'logged_in');assert.equal(result.uid,'123');
 sdk.ctx.option.userId='999';result=await vm.runInNewContext(source,context);assert.equal(result.code,'ACCOUNT_IDENTITY_MISMATCH');
 context.window={};result=await vm.runInNewContext(source,context);assert.equal(result.status,'manual_login');assert.equal(result.reason,'awaiting_login');
})().catch(e=>{console.error(e);process.exitCode=1});
'''.replace('SOURCE',json.dumps(private_identity_expression('123')))
        result=subprocess.run(['node','-e',script],capture_output=True,text=True,timeout=10)
        self.assertEqual(result.returncode,0,result.stderr)

    def test_private_message_entry_reuses_logged_in_profile_without_qr(self):
        from .platform.douyin import CHAT_PAGE_URL, CHAT_ENTRY_SCRIPT
        self.assertEqual(CHAT_PAGE_URL,'https://www.douyin.com/user/self')
        self.assertIn('im-entry',CHAT_ENTRY_SCRIPT)
        self.assertIn('something-button',CHAT_ENTRY_SCRIPT)
        self.assertIn('imSaasContainerId',CHAT_ENTRY_SCRIPT)
        self.assertNotIn('get_qrcode',CHAT_ENTRY_SCRIPT)
        script='''
const vm=require('node:vm'),assert=require('node:assert/strict');let opened=false,clicks=0;
const container={getBoundingClientRect:()=>({width:opened?700:0,height:opened?400:0})};
const entry={click:()=>{clicks++;opened=true}};
const document={querySelector:selector=>selector==='#imSaasContainerId'?container:entry};
(async()=>{const context={document,location:{origin:'https://www.douyin.com'},Date,setTimeout};
assert.equal((await vm.runInNewContext(SOURCE,context)).ready,true);
assert.equal((await vm.runInNewContext(SOURCE,context)).ready,true);assert.equal(clicks,1);
})().catch(e=>{console.error(e);process.exitCode=1});
'''.replace('SOURCE',json.dumps(CHAT_ENTRY_SCRIPT))
        result=subprocess.run(['node','-e',script],capture_output=True,text=True,timeout=10)
        self.assertEqual(result.returncode,0,result.stderr)

    def test_gateway_requests_queue_before_the_native_exclusive_lock(self):
        active=threading.Lock();entered=threading.Event();release=threading.Event();second=threading.Event();errors=[]
        @contextmanager
        def native_lock(alias):
            if not active.acquire(blocking=False):raise AssertionError('native lock overlapped')
            try:yield
            finally:active.release()
        runtimes=Mock();runtimes.alias_lock.side_effect=native_lock
        gateway=Gateway(runtimes,'token','node',Mock())
        def first():
            try:
                with gateway._alias_lock('account'):entered.set();release.wait(3)
            except BaseException as e:errors.append(e)
        def following():
            try:
                with gateway._alias_lock('account'):second.set()
            except BaseException as e:errors.append(e)
        t1=threading.Thread(target=first);t2=threading.Thread(target=following);t1.start();self.assertTrue(entered.wait(1));t2.start()
        try:self.assertFalse(second.wait(.1));self.assertEqual(errors,[])
        finally:release.set();t1.join(3);t2.join(3)
        self.assertTrue(second.is_set());self.assertEqual(errors,[])

    def test_extra_private_login_routes_are_removed(self):
        handler=object.__new__(GatewayHandler);gateway=Mock();handler.server_as_gateway=Mock(return_value=Mock(gateway=gateway))
        for action in ['private-login-qr','private-login-verify']:
            with self.assertRaises(RequestError):handler._route('POST','/v1/browsers/account/douyin/'+action,{}, {'expected_uid':'123'})

    def test_listener_pacing_does_not_wait_while_holding_alias_lock(self):
        gateway = object.__new__(Gateway)
        gateway._alias_lock = lambda alias: threading.Lock()
        gateway._require_douyin_generation = Mock()
        gateway.subscriptions = Mock()
        gateway.subscriptions.poll.return_value = []
        payload = {'runtime_id':'a'*64, 'network_id':'native-'+'b'*32, 'binding_version':1}
        self.assertEqual(gateway.poll_douyin_events('account', payload, {'wait':['25']}), [])
        gateway.subscriptions.poll.assert_called_once_with('account', 50, 0)

    def test_listener_restart_validates_checkpoint_before_touching_browser(self):
        gateway = object.__new__(Gateway)
        gateway.subscriptions = Mock()
        with self.assertRaises(RequestError):
            gateway.start_douyin_events('account', {'expected_uid':'123','notice_checkpoints':{'960':'1'}})
        gateway.subscriptions.start.assert_not_called()

    def test_inbox_route_is_dispatched(self):
        handler = object.__new__(GatewayHandler)
        gateway = Mock()
        server = Mock(gateway=gateway)
        handler.server_as_gateway = Mock(return_value=server)
        body = {'expected_uid':'123'}
        handler._route('POST', '/v1/browsers/account/douyin/inbox', {}, body)
        gateway.douyin_inbox.assert_called_once_with('account',body)

    def test_inbox_script_uses_identity_and_server_messages(self):
        script = inbox_expression('123')
        for text in ['getConversationList', 'getMessageList', 's:server_message_create_time', 's:sender_uid', 'serverId', 'ACCOUNT_IDENTITY_MISMATCH', 'IM_INBOX_LIMIT_EXCEEDED']:
            self.assertIn(text, script)
        self.assertNotIn('getMessagesByConversation', script)
        self.assertNotIn('sendMessage(', script)
        self.assertNotIn('Date.now()', script)
        self.assertIn('123', script)

    def test_sdk_network_error_does_not_claim_definitive_rejection(self):
        script='''
const vm=require('node:vm'),assert=require('node:assert/strict');
const source=SOURCE;
let result={success:false,statusCode:1008,status:1};
const conversation={type:1,id:'conversation',shortId:'short',toParticipantUserId:'456'};
const message={serverId:'100',clientId:'client',sender:'123',type:7,content:'{}',sendFunc(){}};
const sdk={ctx:{option:{userId:'123'}},getConversationList:()=>[conversation],createMessage:async()=>message,sendMessage:async()=>result};
const req={m:{one:function(){/* getOrCreatePrivateConversationByUid */}},c:{one:{exports:{A:{instance:{imSdkService:{imSdkManager:{getImSdkInstance:()=>sdk}}}}}}}};
const chunks=[];chunks.push=function(args){args[2](req)};
const context={location:{origin:'https://www.douyin.com'},fetch:async()=>({ok:true,json:async()=>({status_code:0,user:{uid:'123'}})}),AbortSignal:{timeout:()=>null},Date,setTimeout:()=>0,window:{'@pc-im/im:version':chunks}};
(async()=>{
 let response=await vm.runInNewContext(source,context);assert.equal(response.status,'unknown');assert.equal(response.code,'SDK_NETWORK_ERROR');assert.equal(response.status_code,1008);
 result={success:false,statusCode:0,checkCode:5,checkMessage:'unavailable'};response=await vm.runInNewContext(source,context);assert.equal(response.status,'failed');assert.equal(response.check_code,'5');
 result={success:false,statusCode:0};response=await vm.runInNewContext(source,context);assert.equal(response.status,'unknown');
})().catch(e=>{console.error(e);process.exitCode=1});
'''.replace('SOURCE',json.dumps(im_expression({'uid':'456','action':'send','text':'你好','confirm':True},'123')))
        result=subprocess.run(['node','-e',script],capture_output=True,text=True,timeout=10)
        self.assertEqual(result.returncode,0,result.stderr)

    def test_inbox_reads_live_sdk_cache_and_honors_committed_checkpoints(self):
        script = '''
const vm=require('node:vm'),assert=require('node:assert/strict');
const source=SOURCE;
let seen=0;
const message={type:7,serverId:'100',sender:'456',content:'{"aweType":700,"text":"你好"}',ext:{'s:server_message_create_time':'1700000000000'}};
const conversation={type:1,toParticipantUserId:'456',lastMessage:message,coreInfo:{name:'好友'},getMessageList(){seen++;return [message]}};
const sdk={ctx:{option:{userId:'123'}},getConversationList:()=>[conversation]};
const req={m:{one:function(){/* getOrCreatePrivateConversationByUid */}},c:{one:{exports:{A:{instance:{imSdkService:{imSdkManager:{getImSdkInstance:()=>sdk}}}}}}}};
const chunks=[];chunks.push=function(args){args[2](req)};
const context={location:{origin:'https://www.douyin.com'},fetch:async()=>{throw Error('inbox polling must not make repeated profile requests')},AbortSignal:{timeout:()=>null},crypto:{randomUUID:()=> 'id'},window:{'@pc-im/im:version':chunks}};
(async()=>{
 let result=await vm.runInNewContext(source,context);
 assert.equal(result.status,'succeeded');assert.equal(result.messages[0].server_id,'100');assert.equal(result.checkpoints['456'],'100');assert.equal(seen,1);
 result=await vm.runInNewContext(source.replace('const expected="123", checkpoints={}', 'const expected="123", checkpoints={"456":"100"}'),context);
 assert.equal(result.messages.length,0);assert.equal(seen,1);
 sdk.ctx.option.userId='999';result=await vm.runInNewContext(source,context);assert.equal(result.code,'ACCOUNT_IDENTITY_MISMATCH');sdk.ctx.option.userId='123';
 message.sender='777';result=await vm.runInNewContext(source,context);assert.equal(result.code,'MESSAGE_IDENTITY_INVALID');
 conversation.toParticipantUserId='bad';result=await vm.runInNewContext(source,context);assert.equal(result.code,'CONVERSATION_IDENTITY_INVALID');
 context.window={};result=await vm.runInNewContext(source,context);assert.equal(result.code,'IM_SDK_NOT_READY');
})().catch(e=>{console.error(e);process.exitCode=1});
'''.replace('SOURCE',json.dumps(inbox_expression('123')))
        result=subprocess.run(['node','-e',script],capture_output=True,text=True,timeout=10)
        self.assertEqual(result.returncode,0,result.stderr)

    def test_action_ownership_write_returns_a_cdp_value(self):
        browser=object.__new__(DouyinBrowser)
        browser._evaluate=Mock()
        browser.set_action_ownership('account',{'operation_id':'one'})
        self.assertIn('return true',browser._evaluate.call_args.args[1])

    def test_manual_dm_uses_the_chat_page_not_the_primary_page(self):
        browser=DouyinBrowser()
        private=Mock();private._evaluate.return_value={'status':'logged_in','uid':'123'};private._confirmed_evaluate.return_value={'status':'succeeded'}
        browser._chat_browser=Mock(return_value=private)
        browser.identity=Mock(side_effect=AssertionError('primary page used for DM'))
        result=browser.action('account','123','dm','456',text='你好',confirm=True)
        self.assertEqual(result['status'],'succeeded')
        private.identity.assert_not_called()
        private._evaluate.assert_called_once_with('account',private_identity_expression('123'))
        private._confirmed_evaluate.assert_called_once()
        browser.identity.assert_not_called()

    def test_chat_page_is_created_once_and_never_navigates_the_primary_page(self):
        browser=DouyinBrowser(endpoint=lambda alias:'http://127.0.0.1:19000')
        from urllib.parse import urlsplit
        primary={'id':'primary','type':'page','url':'https://www.douyin.com/user/self'}
        chat={'id':'chat','type':'page','url':'https://www.douyin.com/user/self#__creatorhub_private_messages'}
        browser._target_list=Mock(side_effect=[(urlsplit('http://127.0.0.1:19000'),19000,[primary]),(urlsplit('http://127.0.0.1:19000'),19000,[primary,chat])])
        browser._create_chat_target=Mock(return_value='chat')
        dedicated=Mock();cdp=Mock();cdp.evaluate.side_effect=lambda expression:'/user/self' if expression=='location.pathname' else {'ready':True}
        dedicated.connection=lambda alias:nullcontext(cdp);browser._browser_for_target=Mock(return_value=dedicated)
        self.assertIs(browser._chat_browser('account'),dedicated)
        self.assertIs(browser._chat_browser('account'),dedicated)
        browser._create_chat_target.assert_called_once()
        browser._browser_for_target.assert_called_with('chat')

    def test_chat_page_is_rediscovered_after_gateway_restart(self):
        browser=DouyinBrowser(endpoint=lambda alias:'http://127.0.0.1:19000')
        from urllib.parse import urlsplit
        browser._target_list=Mock(return_value=(urlsplit('http://127.0.0.1:19000'),19000,[{'id':'primary','type':'page','url':'https://www.douyin.com/user/self'},{'id':'chat','type':'page','url':'https://www.douyin.com/user/self#__creatorhub_private_messages'}]))
        browser._is_chat_target=Mock(side_effect=lambda alias,target:target['id']=='chat')
        browser._create_chat_target=Mock()
        dedicated=Mock();cdp=Mock();cdp.evaluate.side_effect=lambda expression:'/user/self' if expression=='location.pathname' else {'ready':True}
        dedicated.connection=lambda alias:nullcontext(cdp);browser._browser_for_target=Mock(return_value=dedicated)
        self.assertIs(browser._chat_browser('account'),dedicated)
        browser._create_chat_target.assert_not_called()
        self.assertEqual(browser._chat_targets['account'],'chat')

    def test_chat_entry_does_not_wait_on_background_page_timers(self):
        from .platform.douyin import CHAT_ENTRY_SCRIPT
        script = '''
const vm=require('node:vm'),assert=require('node:assert/strict');let clicks=0;
const context={document:{querySelector:s=>s==='#imSaasContainerId'?null:{click:()=>clicks++}},location:{origin:'https://www.douyin.com'},Date,setTimeout:()=>{throw Error('background timers cannot drive readiness')}};
const result=vm.runInNewContext(SOURCE,context);
assert.equal(typeof result.then,'undefined');assert.equal(result.ready,false);assert.equal(result.clicked,true);assert.equal(clicks,1);
'''.replace('SOURCE', json.dumps(CHAT_ENTRY_SCRIPT))
        result = subprocess.run(['node', '-e', script], capture_output=True, text=True, timeout=10)
        self.assertEqual(result.returncode, 0, result.stderr)

    def chat_setup(self, cdp):
        from urllib.parse import urlsplit
        browser = DouyinBrowser(endpoint=lambda alias: 'http://127.0.0.1:19000')
        browser._target_list = Mock(return_value=(urlsplit('http://127.0.0.1:19000'), 19000, [{'id': 'chat', 'type': 'page', 'url': 'https://www.douyin.com/user/self#__creatorhub_private_messages'}]))
        browser._chat_targets['account'] = 'chat'
        dedicated = Mock()
        dedicated.connection = lambda alias: nullcontext(cdp)
        browser._browser_for_target = Mock(return_value=dedicated)
        return browser, dedicated

    def test_stale_chat_page_is_refreshed_once_without_touching_the_primary_page(self):
        from .platform.douyin import CHAT_ENTRY_SCRIPT
        cdp = Mock(); refreshed = False
        def evaluate(script):
            if script == 'location.pathname': return '/user/self'
            return {'ready': refreshed, 'clicked': not refreshed}
        def call(method, *args, **kwargs):
            nonlocal refreshed
            if method == 'Page.reload': refreshed = True
        cdp.evaluate.side_effect = evaluate; cdp.command.side_effect = call
        browser, dedicated = self.chat_setup(cdp)
        with patch('browser_gateway.platform.douyin.time.monotonic', side_effect=[0, 0, 13, 13, 13]), patch('browser_gateway.platform.douyin.time.sleep'):
            self.assertIs(browser._chat_browser('account'), dedicated)
        self.assertEqual([c.args[0] for c in cdp.command.call_args_list], ['Page.enable', 'Page.reload'])
        cdp.wait_event.assert_called_once()
        self.assertEqual(cdp.wait_event.call_args.args[0], 'Page.domContentEventFired')
        self.assertEqual(sum(c.args[0] == CHAT_ENTRY_SCRIPT for c in cdp.evaluate.call_args_list), 2)
        browser._browser_for_target.assert_called_once_with('chat')

    def test_chat_polling_clicks_entry_once_and_still_reports_failure_after_refresh(self):
        from .platform.douyin import CHAT_ENTRY_SCRIPT
        cdp = Mock()
        cdp.evaluate.side_effect = lambda script: '/user/self' if script == 'location.pathname' else {'ready': False, 'clicked': True}
        browser, _ = self.chat_setup(cdp)
        with patch('browser_gateway.platform.douyin.time.monotonic', side_effect=[0, 0, 1, 13, 13, 13, 14, 26]), patch('browser_gateway.platform.douyin.time.sleep'):
            with self.assertRaisesRegex(DouyinError, 'private-message panel.*refresh'):
                browser._chat_browser('account')
        self.assertEqual(sum(c.args[0] == CHAT_ENTRY_SCRIPT for c in cdp.evaluate.call_args_list), 2)
        self.assertEqual([c.args[0] for c in cdp.command.call_args_list], ['Page.enable', 'Page.reload'])

    def test_chat_invalid_response_is_not_hidden_by_refresh(self):
        cdp = Mock()
        cdp.evaluate.side_effect = ['/user/self', {}]
        browser, _ = self.chat_setup(cdp)
        with self.assertRaisesRegex(DouyinError, 'response is invalid'):
            browser._chat_browser('account')
        cdp.command.assert_not_called()

    def test_browser_inbox_validation_and_response(self):
        browser = object.__new__(DouyinBrowser)
        browser.identity = Mock()
        browser._evaluate = Mock(return_value={'status': 'succeeded', 'account_uid': '123', 'messages': []})
        browser._chat_browser = Mock(return_value=browser)
        self.assertEqual(browser.inbox('account', '123')['status'], 'succeeded')
        browser.identity.assert_not_called()
        with self.assertRaises(DouyinError): browser.inbox('account', 'invalid')
        browser._evaluate.return_value = []
        with self.assertRaises(DouyinError): browser.inbox('account', '123')

    def test_gateway_generation_and_response(self):
        gateway = object.__new__(Gateway)
        gateway.browser = Mock()
        gateway.browser.inbox.return_value = {'status':'succeeded','account_uid':'123','messages':[]}
        gateway._alias_lock = Mock(return_value=nullcontext())
        gateway._require_douyin_generation = Mock()
        payload = {'runtime_id':'a'*64,'network_id':'native-'+'b'*32,'binding_version':1,'expected_uid':'123'}
        self.assertEqual(gateway.douyin_inbox('alias', payload)['status'],'succeeded')
        self.assertEqual(gateway._require_douyin_generation.call_count,2)
        gateway.browser.inbox.assert_called_once_with('alias','123',{})
        for bad in [{},dict(payload,expected_uid='bad'),dict(payload,binding_version=0)]:
            with self.assertRaises(RequestError):gateway.douyin_inbox('alias',bad)
        gateway.browser.inbox.return_value = {'status':'failed','code':'IM_SDK_NOT_READY'}
        self.assertEqual(gateway.douyin_inbox('alias',payload)['code'],'IM_SDK_NOT_READY')
        gateway.browser.inbox.side_effect = DouyinError('broken')
        with self.assertRaises(RequestError):gateway.douyin_inbox('alias',payload)