| |
| """Regression and local-media experiment for adaptive handling of TikTok format changes.""" |
| import collections |
| import json |
| import os |
| import pathlib |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import threading |
| import time |
| import unittest |
| from unittest.mock import patch |
|
|
| ROOT=pathlib.Path(sys.argv[1]).resolve() if len(sys.argv)>1 else pathlib.Path(__file__).parent |
| sys.path.insert(0,str(ROOT)) |
| from flv_bridge import FLVBridge |
| from ingest import RelayHub |
| from relay import Sender |
| from youtube_manager import ManagedYouTube |
|
|
|
|
| def tag(kind,ts,body): |
| size=len(body) |
| header=bytes([kind])+size.to_bytes(3,'big')+bytes([ |
| (ts>>16)&255,(ts>>8)&255,ts&255,(ts>>24)&255,0,0,0]) |
| return header+body+(11+size).to_bytes(4,'big') |
|
|
|
|
| def flv(*records): |
| return b'FLV\x01\x05\x00\x00\x00\x09\x00\x00\x00\x00'+b''.join(records) |
|
|
|
|
| def seq(cfg):return b'\x17\x00\x00\x00\x00'+cfg |
|
|
|
|
| def coded(payload=b'\x00\x00\x00\x02\x65\x88'): |
| return b'\x27\x01\x00\x00\x00'+payload |
|
|
|
|
| AVC_A=bytes.fromhex('0142c00dffe100166742c00dda0507ec0440000003004000000ca3c50aa801000468ce0fc8') |
| AVC_B=bytes.fromhex('0142c01effe100196742c01eda0280bfe5c044000003000400000300ca3c58ba8001000468ce0fc8') |
|
|
|
|
| class FakePipe: |
| def __init__(self):self.close_count=0 |
| def close(self):self.close_count+=1 |
|
|
|
|
| class FakeProcess: |
| def __init__(self,exit_code=None):self.stdin=FakePipe();self.exit_code=exit_code |
| def poll(self):return self.exit_code |
|
|
|
|
| class AdaptivePolicyTests(unittest.TestCase): |
| def setUp(self): |
| self.env=patch.dict(os.environ,{ |
| 'YOUTUBE_LIGHT_MAX':'2','YOUTUBE_FLAP_LIMIT':'3', |
| 'YOUTUBE_FLAP_WINDOW_SECONDS':'600','YOUTUBE_FLAP_FALLBACK_MODE':'matchgop', |
| 'YOUTUBE_SEND_MODE':'auto','YOUTUBE_AUTO_TRANSCODE':'0'}) |
| self.env.start();self.addCleanup(self.env.stop) |
| self.light_hub=RelayHub() |
| self.hub_patch=patch('ingest.hub',self.light_hub) |
| self.hub_patch.start();self.addCleanup(self.hub_patch.stop) |
|
|
| def make_sender(self): |
| messages=[];sender=Sender('/bin/true',lambda user,text:messages.append((user,text))) |
| sender.current_candidate={'codec':'h264','label':'h264','width':320,'height':240} |
| sender.lock_gold(sender.current_candidate) |
| sender.active_mode='copy' |
| sender.session={'force_stop':False} |
| sender.process=FakeProcess() |
| persisted=[];sender.on_sticky=lambda *args:persisted.append(args) |
| return sender,messages,persisted |
|
|
| def test_real_flv_resolution_events_trigger_one_per_sender_fallback_at_threshold(self): |
| sender,messages,persisted=self.make_sender() |
| with patch('relay.log_event') as saved: |
| bridge=FLVBridge(user='alice',event_logger=sender._bridge_event) |
| bridge.feed(flv(tag(9,0,seq(AVC_A)),tag(9,1,coded()))) |
| self.assertIsNone(sender.sticky_mode) |
| self.assertEqual(sender.process.stdin.close_count,0) |
| for ts,config in ((40,AVC_B),(80,AVC_A)): |
| bridge.feed(tag(9,ts,seq(config))) |
| bridge.feed(tag(9,ts+1,coded())) |
| self.assertIsNone(sender.sticky_mode) |
| self.assertEqual(sender.process.stdin.close_count,0) |
| bridge.feed(tag(9,120,seq(AVC_B))) |
| self.assertEqual(bridge.format_changes,3) |
| self.assertEqual(sender.sticky_mode,'matchgop') |
| self.assertTrue(sender.session['force_stop']) |
| self.assertEqual(sender.process.stdin.close_count,1) |
| self.assertEqual(persisted,[('alice','matchgop',320,240)]) |
| self.assertEqual(sum(c.args[1]=='RESOLUTION_CHANGE' for c in saved.call_args_list),3) |
| self.assertTrue(any(c.args[1]=='FORMAT_FALLBACK' for c in saved.call_args_list)) |
| bridge.feed(tag(9,160,seq(AVC_A))) |
| self.assertEqual(sender.process.stdin.close_count,1,'one format flap must not restart the RTMP publisher repeatedly') |
| self.assertTrue(any('sticky matchgop' in text for _,text in messages)) |
|
|
| def test_full_light_cap_keeps_copy_alive_then_retries_without_overbooking(self): |
| self.assertTrue(self.light_hub.reserve_light('peer-a')) |
| self.assertTrue(self.light_hub.reserve_light('peer-b')) |
| sender,_,_=self.make_sender() |
| bridge=FLVBridge(user='alice',event_logger=sender._bridge_event) |
| with patch('relay.log_event') as saved: |
| bridge.feed(flv(tag(9,0,seq(AVC_A)),tag(9,1,coded()))) |
| for ts,config in ((40,AVC_B),(80,AVC_A),(120,AVC_B)): |
| bridge.feed(tag(9,ts,seq(config))) |
| bridge.feed(tag(9,ts+1,coded())) |
| self.assertIsNone(sender.sticky_mode) |
| self.assertEqual(sender.pending_sticky[0],'matchgop') |
| self.assertFalse(sender.session['force_stop']) |
| self.assertEqual(sender.process.stdin.close_count,0) |
| self.assertEqual(self.light_hub.light_count(),2) |
| self.light_hub.release_light('peer-a') |
| self.assertTrue(sender.apply_pending_sticky('alice')) |
| self.assertEqual(sender.sticky_mode,'matchgop') |
| self.assertIsNone(sender.pending_sticky) |
| self.assertEqual(sender.process.stdin.close_count,1) |
| self.assertEqual(self.light_hub.light_count(),2) |
| self.assertTrue(any(call.args[1]=='FORMAT_FALLBACK' for call in saved.call_args_list)) |
|
|
| def test_format_changes_outside_window_do_not_force_transcoding(self): |
| sender,_,_=self.make_sender() |
| sender.flap.append(time.monotonic()-601) |
| self.assertFalse(sender.note_format_change('alice',{'old_width':320,'old_height':240, |
| 'new_width':640,'new_height':360})) |
| self.assertEqual(len(sender.flap),1) |
| self.assertIsNone(sender.sticky_mode) |
| self.assertEqual(sender.process.stdin.close_count,0) |
|
|
| def test_reservation_is_released_with_relay_owner(self): |
| self.assertTrue(self.light_hub.reserve_light('alice')) |
| self.assertEqual(self.light_hub.light_count(),1) |
| self.light_hub.release('alice','oauth') |
| self.assertEqual(self.light_hub.light_count(),0) |
|
|
| def test_light_reservation_is_atomic_under_simultaneous_fallbacks(self): |
| with patch.dict(os.environ,{'YOUTUBE_LIGHT_MAX':'1'}): |
| gate=threading.Barrier(3);results=[] |
| def reserve(user): |
| gate.wait();results.append(self.light_hub.reserve_light(user)) |
| threads=[threading.Thread(target=reserve,args=(u,)) for u in ('alice','bob')] |
| for thread in threads:thread.start() |
| gate.wait() |
| for thread in threads:thread.join(timeout=2) |
| self.assertEqual(sum(bool(x) for x in results),1) |
| self.assertEqual(self.light_hub.light_count(),1) |
|
|
| def test_sender_snapshot_reports_effective_mode_not_ui_preference(self): |
| sender=Sender('/bin/true',lambda *a:None) |
| sender.mode_override='auto' |
| sender.active_mode='matchgop' |
| self.assertEqual(sender.snapshot()['send_mode'],'matchgop') |
|
|
| def test_output_failure_reason_is_not_mislabeled_as_normal_stop(self): |
| sender=Sender('/bin/true',lambda *a:None) |
| sender.last_failure_side='output' |
| reason=sender._close_reason({},FakeProcess(exit_code=187)) |
| self.assertIn('RTMPS',reason) |
| self.assertIn('187',reason) |
| self.assertNotEqual(reason,'إيقاف عادي') |
|
|
| def test_explicit_stop_keeps_normal_stop_diagnostic(self): |
| sender=Sender('/bin/true',lambda *a:None) |
| self.assertEqual(sender._close_reason({},FakeProcess(exit_code=0)),'إيقاف عادي') |
|
|
|
|
| class EventModePersistenceTests(unittest.TestCase): |
| def setUp(self): |
| self.tmp=tempfile.TemporaryDirectory(prefix='sender-mode-journal-') |
| self.addCleanup(self.tmp.cleanup) |
| self.path=pathlib.Path(self.tmp.name)/'journal.json' |
| self.manager=ManagedYouTube('/bin/true',lambda u:None,lambda *a:None,lambda:['alice'],self.path) |
| self.key='alice:room' |
| event={'username':'alice','session_key':self.key,'room_id':'room','broadcast_id':'VIDEO', |
| 'title':'test','nickname':'Alice','phase':'live','youtube_status':'live','closing':False} |
| self.manager.journal.change(lambda d:d['events'].update({self.key:event})) |
|
|
| def test_adaptive_mode_canvas_and_pending_request_survive_manager_recreation(self): |
| self.manager._remember_pending_sender_mode('alice','matchgop','four source resolution changes') |
| original=self.manager.journal.read()['events'][self.key] |
| self.assertEqual(original['sender_pending_mode'],'matchgop') |
| restored=ManagedYouTube('/bin/true',lambda u:None,lambda *a:None,lambda:['alice'],self.path) |
| restored._remember_sender_mode('alice','matchgop',720,1280) |
| ev=restored.journal.read()['events'][self.key] |
| self.assertEqual((ev['sender_mode'],ev['sender_canvas_width'],ev['sender_canvas_height']), |
| ('matchgop',720,1280)) |
| self.assertEqual(ev['sender_pending_mode'],'') |
|
|
| def test_journaled_mode_is_visible_before_sender_thread_restarts(self): |
| self.manager._remember_sender_mode('alice','matchgop',720,1280) |
| self.manager._remember_pending_sender_mode('alice','matchgop','format flapping') |
| row=self.manager.snapshot()['events'][0] |
| self.assertEqual(row['send_mode'],'matchgop') |
| self.assertEqual(row['pending_send_mode'],'matchgop') |
|
|
| def test_archiving_event_clears_per_event_mode_and_pending_reservation(self): |
| self.manager._remember_sender_mode('alice','matchgop',720,1280) |
| self.manager._remember_pending_sender_mode('alice','matchgop','format flapping') |
| self.manager._sender_sticky['alice']='matchgop' |
| self.manager._sender_gop['alice']=(720,1280) |
| self.manager._sender_pending['alice']=('matchgop','format flapping') |
| self.manager.archive_event(self.key,'done') |
| self.assertNotIn(self.key,self.manager.journal.read()['events']) |
| self.assertNotIn('alice',self.manager._sender_sticky) |
| self.assertNotIn('alice',self.manager._sender_gop) |
| self.assertNotIn('alice',self.manager._sender_pending) |
|
|
|
|
| class IdleHousekeepingTests(unittest.TestCase): |
| def setUp(self): |
| self.tmp=tempfile.TemporaryDirectory(prefix='idle-clean-') |
| self.addCleanup(self.tmp.cleanup) |
| self.path=pathlib.Path(self.tmp.name)/'journal.json' |
| self.manager=ManagedYouTube('/bin/true',lambda u:None,lambda *a:None,lambda:['alice'],self.path) |
| self.manager.idle_reset_seconds=600.0 |
| self.manager.journal.change(lambda d:(d['streams'].update({'alice':{'id':'STREAM1'}}), |
| d.update(channel={'id':'CHAN','title':'قناتي'}))) |
| self.manager._sender_sticky['alice']='matchgop' |
| self.manager._sender_gop['alice']=(720,1280) |
| self.manager.hammer['alice']=3 |
| self.manager.destinations['alice']='rtmps://cached.example.invalid/live/KEY' |
|
|
| def test_no_clean_before_full_idle_threshold_then_clean_once(self): |
| m=self.manager |
| self.assertFalse(m.idle_housekeeping(now=1000.0)) |
| self.assertFalse(m.idle_housekeeping(now=1599.0)) |
| self.assertEqual(m._sender_sticky.get('alice'),'matchgop') |
| self.assertTrue(m.idle_housekeeping(now=1600.0)) |
| self.assertEqual(m._sender_sticky,{}) |
| self.assertEqual(m._sender_gop,{}) |
| self.assertEqual(m.hammer,{}) |
| self.assertEqual(m.destinations,{}) |
| self.assertFalse(m.idle_housekeeping(now=9999.0)) |
|
|
| def test_clean_preserves_channel_binding_streams_and_settings(self): |
| m=self.manager |
| m.idle_housekeeping(now=0.0);m.idle_housekeeping(now=600.0) |
| d=m.journal.read() |
| self.assertEqual(d['channel'],{'id':'CHAN','title':'قناتي'}) |
| self.assertEqual(d['streams']['alice']['id'],'STREAM1') |
| self.assertIn('title',d['settings']) |
| self.assertGreater(d.get('last_idle_clean') or 0,0) |
|
|
| def test_any_event_counts_as_activity_and_restarts_the_timer(self): |
| m=self.manager |
| key='alice:room' |
| event={'username':'alice','session_key':key,'room_id':'room','broadcast_id':'VIDEO', |
| 'title':'test','nickname':'Alice','phase':'live','youtube_status':'live','closing':False} |
| m.journal.change(lambda d:d['events'].update({key:event})) |
| self.assertFalse(m.idle_housekeeping(now=1000.0)) |
| self.assertFalse(m.idle_housekeeping(now=99999.0)) |
| self.assertEqual(m._sender_sticky.get('alice'),'matchgop') |
| m.journal.change(lambda d:d['events'].pop(key,None)) |
| self.assertFalse(m.idle_housekeeping(now=100000.0)) |
| self.assertFalse(m.idle_housekeeping(now=100599.0)) |
| self.assertTrue(m.idle_housekeeping(now=100600.0)) |
|
|
| def test_zero_disables_housekeeping(self): |
| self.manager.idle_reset_seconds=0.0 |
| self.assertFalse(self.manager.idle_housekeeping(now=10**9)) |
| self.assertEqual(self.manager._sender_sticky.get('alice'),'matchgop') |
|
|
|
|
| @unittest.skipUnless(shutil.which('ffmpeg') and shutil.which('ffprobe'),'FFmpeg and ffprobe required') |
| class RealResolutionTransitionTests(unittest.TestCase): |
| def setUp(self): |
| self.tmp=tempfile.TemporaryDirectory(prefix='adaptive-format-') |
| self.addCleanup(self.tmp.cleanup) |
| self.root=pathlib.Path(self.tmp.name) |
| self.ff=shutil.which('ffmpeg');self.fp=shutil.which('ffprobe') |
|
|
| def run_checked(self,args,**kw): |
| return subprocess.run(args,check=True,capture_output=True,timeout=60,**kw) |
|
|
| def probe(self,path,*extra): |
| return json.loads(self.run_checked([self.fp,'-v','error',*extra,'-of','json',str(path)],text=True).stdout) |
|
|
| def test_copy_preserves_dynamic_dimensions_and_matchgop_normalizes_them(self): |
| sources=[] |
| for n,(width,height) in enumerate(((320,640),(480,960))): |
| path=self.root/('source-%d.flv'%n) |
| self.run_checked([self.ff,'-hide_banner','-v','error','-y','-f','lavfi','-i', |
| 'testsrc2=size=%dx%d:rate=20'%(width,height),'-f','lavfi','-i','sine=sample_rate=48000', |
| '-t','2.5','-c:v','libx264','-preset','ultrafast','-pix_fmt','yuv420p', |
| '-g','40','-keyint_min','40','-sc_threshold','0','-c:a','aac','-b:a','96k', |
| '-ar','48000','-f','flv',str(path)]) |
| sources.append(path.read_bytes()) |
| observed=[] |
| bridge=FLVBridge(user='fixture',event_logger=lambda user,event,details:observed.append((event,details))) |
| joined=bridge.feed(sources[0]);bridge.reconnect();joined+=bridge.feed(sources[1]) |
| source=self.root/'transition.flv';source.write_bytes(joined) |
| self.assertEqual(bridge.format_changes,1) |
| self.assertTrue(any(event=='RESOLUTION_CHANGE' for event,_ in observed)) |
|
|
| copied=self.root/'copy.flv' |
| self.run_checked([self.ff,'-hide_banner','-v','error','-y','-i',str(source), |
| '-map','0:v:0','-map','0:a:0?','-c:v','copy','-c:a','aac','-ar','48000','-f','flv',str(copied)]) |
| normalized=self.root/'matchgop.flv' |
| vf=('scale=w=320:h=640:force_original_aspect_ratio=decrease:flags=lanczos:eval=frame,' |
| 'pad=320:640:(ow-iw)/2:(oh-ih)/2:color=black:eval=frame,setsar=1,format=yuv420p') |
| self.run_checked([self.ff,'-hide_banner','-v','error','-y','-i',str(source), |
| '-map','0:v:0','-map','0:a:0?','-vf',vf,'-c:v','libx264','-preset','ultrafast', |
| '-tune','zerolatency','-b:v','900k','-minrate','900k','-maxrate','900k','-bufsize','1350k', |
| '-force_key_frames','expr:gte(t,n_forced*2)','-g','40','-keyint_min','30','-sc_threshold','0', |
| '-bf','0','-x264-params','nal-hrd=cbr:force-cfr=1:open-gop=0:scenecut=0:keyint=40:min-keyint=30', |
| '-pix_fmt','yuv420p','-fps_mode','cfr','-c:a','aac','-ar','48000','-f','flv',str(normalized)]) |
|
|
| copied_frames=self.probe(copied,'-select_streams','v:0','-show_frames','-show_entries','frame=width,height')['frames'] |
| copied_sizes=collections.Counter((int(f['width']),int(f['height'])) for f in copied_frames if 'width' in f) |
| self.assertGreater(copied_sizes[(320,640)],30) |
| self.assertGreater(copied_sizes[(480,960)],30) |
| normalized_frames=self.probe(normalized,'-select_streams','v:0','-show_frames','-show_entries','frame=width,height')['frames'] |
| normalized_sizes=collections.Counter((int(f['width']),int(f['height'])) for f in normalized_frames if 'width' in f) |
| self.assertGreater(len(normalized_frames),90) |
| self.assertEqual(set(normalized_sizes),{(320,640)}) |
|
|
| for output in (copied,normalized): |
| packets=self.probe(output,'-show_packets','-show_entries','packet=stream_index,dts_time')['packets'] |
| per_stream={} |
| for packet in packets: |
| if packet.get('dts_time') is not None: |
| per_stream.setdefault(packet['stream_index'],[]).append(float(packet['dts_time'])) |
| self.assertGreaterEqual(len(per_stream),2) |
| for times in per_stream.values(): |
| self.assertTrue(all(b>=a for a,b in zip(times,times[1:])), |
| 'DTS must not regress at a same-event source handover') |
| self.run_checked([self.ff,'-hide_banner','-v','error','-i',str(output),'-map','0:v:0','-f','null','-']) |
| fmt=self.probe(output,'-show_format')['format'] |
| self.assertGreater(float(fmt['duration']),4.5) |
|
|
|
|
| if __name__=='__main__':unittest.main(argv=[sys.argv[0]],verbosity=2) |
|
|
