Blame for test_adaptive_format.py

#!/usr/bin/env python3
"""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)