Blame for test_source_probe.py

#!/usr/bin/env python3
"""Deterministic FLV pre-start metadata tests; no network or FFmpeg required."""
import os
import pathlib
import sys
import tempfile

ROOT=pathlib.Path(sys.argv[1] if len(sys.argv)>1 else pathlib.Path(__file__).parent).resolve()
sys.path.insert(0,str(ROOT))
os.environ.setdefault('PASSWORD','test-pass-not-for-prod')
os.environ.setdefault('DATA_DIR',tempfile.mkdtemp(prefix='source-probe-test-'))
os.environ['YOUTUBE_FLV_PROBE_SECONDS']='0.2'

import youtube_manager as manager_module  # noqa: E402
from relay import Sender  # noqa: E402
from stream_monitor import video_dimensions  # noqa: E402

passed=failed=0
def check(name,ok,detail=''):
    global passed,failed
    if ok:
        passed+=1;print('PASS ',name)
    else:
        failed+=1;print('FAIL ',name,detail)


class Bits:
    def __init__(self):self.data=[]
    def bit(self,value):self.data.append(1 if value else 0)
    def bits(self,value,count):
        for shift in range(count-1,-1,-1):self.bit((value>>shift)&1)
    def ue(self,value):
        code=value+1;size=code.bit_length()
        self.bits(0,size-1);self.bits(code,size)
    def rbsp(self):
        self.bit(1)
        while len(self.data)%8:self.bit(0)
        raw=bytes(sum(self.data[i+j]<<(7-j) for j in range(8)) for i in range(0,len(self.data),8))
        out=bytearray();zeros=0
        for byte in raw:
            if zeros>=2 and byte<=3:
                out.append(3);zeros=0
            out.append(byte);zeros=zeros+1 if byte==0 else 0
        return bytes(out)


def avcc(width,height):
    """Minimal baseline H.264 SPS/PPS with uncropped, frame-only 4:2:0 dimensions."""
    assert width%16==0 and height%16==0
    b=Bits();b.bits(66,8);b.bits(0,8);b.bits(30,8)  # profile/constraints/level
    b.ue(0);b.ue(0);b.ue(0);b.ue(0);b.ue(1);b.bit(0)
    b.ue(width//16-1);b.ue(height//16-1)
    b.bit(1);b.bit(1);b.bit(0);b.bit(0)
    sps=b'\x67'+b.rbsp();pps=bytes.fromhex('68ce0fc8')
    return (bytes([1,66,0,30,0xff,0xe1])+len(sps).to_bytes(2,'big')+sps+
            b'\x01'+len(pps).to_bytes(2,'big')+pps)


def tag(config,ts=0):
    body=b'\x17\x00\x00\x00\x00'+config;size=len(body)
    header=(b'\x09'+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(*tags):
    return b'FLV\x01\x05\x00\x00\x00\x09\x00\x00\x00\x00'+b''.join(tags)


class Response:
    def __init__(self,chunks):self.chunks=list(chunks);self.closed=False
    def raise_for_status(self):return None
    def iter_content(self):yield from self.chunks
    def close(self):self.closed=True


def manager():
    obj=manager_module.ManagedYouTube.__new__(manager_module.ManagedYouTube)
    obj.headers_for=lambda user:{}
    obj._unstable={};obj.starve={};obj._skip_probe_once={};obj._probe_note=''
    obj.logs=[];obj.log=lambda user,text:obj.logs.append(str(text))
    return obj


def candidate(width=432,height=864):
    return {'url':'https://media.example.test/live.flv','media':'flv','codec':'h264',
            'quality_name':'hd','label':'hd','brief':'hd H264','width':width,'height':height}


def probe_with(chunks,cand=None):
    obj=manager();responses=[]
    original=manager_module.requests.get
    def fake_get(*args,**kwargs):
        response=Response(chunks);responses.append(response);return response
    manager_module.requests.get=fake_get
    try:result=obj.probe('alice',cand or candidate())
    finally:manager_module.requests.get=original
    return result,responses


small=avcc(16,16);real=avcc(640,1280)
check('test SPS generator carries exact dimensions',
      video_dimensions('avc1',small)==(16,16) and video_dimensions('avc1',real)==(640,1280))

blob=flv(tag(real));chunks=[blob[i:i+7] for i in range(0,len(blob),7)]
result,responses=probe_with(chunks)
check('fragmented FLV probe returns actual decoder dimensions',
      isinstance(result,dict) and (result.get('width'),result.get('height'))==(640,1280),str(result))
check('FLV probe closes the source response',bool(responses) and responses[0].closed)

blob=flv(tag(small),tag(real));split=13+len(tag(small))
result,_=probe_with([blob[:split],blob[split:]])
check('sub-32px startup config is not accepted as the final canvas',
      isinstance(result,dict) and (result.get('width'),result.get('height'))==(640,1280),str(result))

result,responses=probe_with([b'FLV'])
check('signature-only FLV keeps the old truthy fallback',result is True,str(result))
check('signature-only fallback also closes the response',bool(responses) and responses[0].closed)
result,_=probe_with([b'NOT AN FLV'])
check('invalid FLV signature is still rejected',result is False,str(result))

obj=manager();obj.probe=lambda user,c:{'width':640,'height':1280,'codec':'avc1'}
selected=obj._select_source('alice',{'sources':[candidate()]})
check('source selection retains a usable candidate',len(selected)==1,str(selected))
if selected:
    chosen=selected[0]
    check('actual dimensions are attached without rewriting the rendition label',
          (chosen.get('width'),chosen.get('height'))==(432,864) and
          (chosen.get('probed_width'),chosen.get('probed_height'))==(640,1280) and
          obj._source_pin(chosen)==obj._source_pin(candidate()),str(chosen))
    sender=Sender('ffmpeg',lambda *args:None)
    check('matchgop locks the actual HD canvas instead of the stale SD label',
          sender.lock_gold(chosen)==(720,1280),str(sender.lock_gold(chosen)))
    check('a declared/probed mismatch is surfaced in operational logs',
          any('640x1280' in line and '432x864' in line for line in obj.logs),str(obj.logs))

fallback=Sender('ffmpeg',lambda *args:None)
check('without decoder metadata the old declared-canvas fallback is unchanged',
      fallback.lock_gold(candidate())==(432,864),str(fallback.lock_gold(candidate())))

print('source probe tests: %d passed, %d failed'%(passed,failed))
raise SystemExit(1 if failed else 0)