Spaces:
Paused
Paused
| """Completion tests: cap, source rank, title clock, stall watchdog, isolation, no recorder. | |
| Simulation of concurrent relays uses real FFmpeg when present, plus RelayHub. | |
| This is NOT a real TikTok LIVE / YouTube RTMPS integration test. | |
| """ | |
| import os, pathlib, subprocess, sys, tempfile, threading, time, json | |
| from datetime import datetime | |
| from zoneinfo import ZoneInfo | |
| from shutil import which | |
| DIR = pathlib.Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else pathlib.Path(__file__).resolve().parent | |
| os.environ.setdefault('PASSWORD', 'admin') | |
| os.environ.setdefault('DATA_DIR', str(tempfile.mkdtemp())) | |
| os.environ.setdefault('YOUTUBE_SEND_MODE', 'copy') | |
| sys.path.insert(0, str(DIR)) | |
| failed = [] | |
| passed = 0 | |
| def check(name, cond, detail=''): | |
| global passed | |
| if cond: | |
| passed += 1 | |
| print('PASS ', name) | |
| else: | |
| failed.append(name) | |
| print('FAIL ', name, detail) | |
| FF = which(os.environ.get('FFMPEG_PATH') or 'ffmpeg') or which('ffmpeg') | |
| # --- 1. concurrent cap --- | |
| os.environ.pop('MAX_CONCURRENT_RELAYS', None) | |
| os.environ.pop('YOUTUBE_MAX_EVENTS', None) | |
| import importlib | |
| import ingest | |
| importlib.reload(ingest) | |
| check('default MAX_CONCURRENT_RELAYS is 2 (cap, not a guarantee)', ingest.max_concurrent_relays() == 2, str(ingest.max_concurrent_relays())) | |
| os.environ['MAX_CONCURRENT_RELAYS'] = '7' | |
| importlib.reload(ingest) | |
| check('env 7 is accepted as a cap (not a promise of 7 healthy streams)', ingest.max_concurrent_relays() == 7) | |
| os.environ['MAX_CONCURRENT_RELAYS'] = '99' | |
| importlib.reload(ingest) | |
| check('typo 99 is clamped to 8', ingest.max_concurrent_relays() == 8) | |
| os.environ['MAX_CONCURRENT_RELAYS'] = '2' | |
| importlib.reload(ingest) | |
| check('env 2 in force', ingest.max_concurrent_relays() == 2) | |
| # --- 2. title Baghdad session clock, not process start --- | |
| from ingest import format_session_title | |
| when = datetime(2026, 9, 27, 21, 30, tzinfo=ZoneInfo('Asia/Baghdad')) | |
| title, date, tm = format_session_title('StrongArmor', when) | |
| check('title template Baghdad', title == 'StrongArmor | 2026-09-27 21:30', title) | |
| boot = datetime(2026, 9, 27, 1, 0, tzinfo=ZoneInfo('Asia/Baghdad')) | |
| title_boot, _, _ = format_session_title('StrongArmor', boot) | |
| check('title uses session time not a fake docker boot time', title != title_boot) | |
| # --- 3. source ranking: 720x1440 H264 is not "heavy 1440p" --- | |
| # Import rank_sources without booting FastAPI password/ffmpeg if possible. | |
| # app.py raises if no PASSWORD (set) and no ffmpeg. | |
| if not FF: | |
| check('ffmpeg present for remaining tests', False, 'ffmpeg missing') | |
| print('%d passed, %d failed' % (passed, len(failed))) | |
| sys.exit(1) | |
| os.environ['FFMPEG_PATH'] = FF | |
| import app as appmod | |
| from youtube_manager import ManagedYouTube | |
| hd = {'url': 'https://example.invalid/hd.flv', 'media': 'flv', 'codec': 'h264', | |
| 'quality_name': 'hd', 'resolution': '720x1440', 'width': 720, 'height': 1440, | |
| 'bitrate': 1000000, 'label': 'hd · 720x1440 · h264', 'brief': 'hd 1440p H264'} | |
| origin = {'url': 'https://example.invalid/origin.flv', 'media': 'flv', 'codec': 'h265', | |
| 'quality_name': 'origin', 'resolution': '960x1920', 'width': 960, 'height': 1920, | |
| 'bitrate': 1892000, 'label': 'origin · 960x1920 · h265', 'brief': 'origin 1920p HEVC'} | |
| sd = {'url': 'https://example.invalid/sd.flv', 'media': 'flv', 'codec': 'h265', | |
| 'quality_name': 'sd', 'resolution': '576x1152', 'width': 576, 'height': 1152, | |
| 'bitrate': 800000, 'label': 'sd · 576x1152 · h265', 'brief': 'sd 1152p HEVC'} | |
| ranked = appmod.rank_sources({'blocks': []}) | |
| # feed via add path: construct live_core style | |
| payload = json.dumps({'data': { | |
| 'hd': {'main': {'flv': hd['url'], 'sdk_params': json.dumps({'resolution': '720x1440', 'VCodec': 'h264', 'vbitrate': 1000000})}}, | |
| 'origin': {'main': {'flv': origin['url'], 'sdk_params': json.dumps({'resolution': '960x1920', 'VCodec': 'h265', 'vbitrate': 1892000})}}, | |
| 'sd': {'main': {'flv': sd['url'], 'sdk_params': json.dumps({'resolution': '576x1152', 'VCodec': 'h265', 'vbitrate': 800000})}}, | |
| }}) | |
| ranked = appmod.rank_sources({'live_core_sdk_data': {'pull_data': {'stream_data': payload}}}) | |
| check('rank_sources returns three pull URLs', len(ranked) >= 3, str(len(ranked))) | |
| top = ranked[0] if ranked else {} | |
| check('H264 FLV ranks above HEVC origin', 'h264' in str(top.get('codec') or '').lower(), str(top.get('brief'))) | |
| check('top is hd not origin', str(top.get('quality_name') or '').lower() == 'hd', str(top.get('quality_name'))) | |
| os.environ['YOUTUBE_VIDEO_MODE'] = 'original' | |
| os.environ['YOUTUBE_PREFER_H264'] = '1' | |
| os.environ['YOUTUBE_SEND_MODE'] = 'auto' | |
| sel = ManagedYouTube.video_sources(ranked) | |
| check('video_sources prefers H264 when original+prefer', sel and 'h264' in str(sel[0].get('codec') or '').lower(), str((sel or [{}])[0].get('brief'))) | |
| # busy>=2 heavy: 720x1440 must survive (min side 720) | |
| # We cannot easily set hub.active_count without occupying; test _too_heavy via selecting when busy mocked. | |
| # Direct: min(720,1440)=720 < 1440 so not heavy. | |
| def _too_heavy(c): | |
| w=int(c.get('width') or 0); h=int(c.get('height') or 0) | |
| if w and h: | |
| return min(w,h)>=1440 or (w*h)>=(1600*900) | |
| text=str(c.get('label') or c.get('brief') or '') | |
| return any(x in text for x in ('2160','4k','4K')) | |
| check('720x1440 is not treated as heavy 1440p', not _too_heavy(hd)) | |
| check('2560x1440 landscape is heavy', _too_heavy({'width':2560,'height':1440})) | |
| check('brief text 1440p alone without dims is not auto-heavy', not _too_heavy({'brief':'hd 1440p H264','width':0,'height':0})) | |
| # --- 4. no recorder in executable path --- | |
| src = (DIR/'app.py').read_text(encoding='utf-8') + (DIR/'relay.py').read_text(encoding='utf-8') | |
| check('no def record(', 'def record(' not in src) | |
| check('no FLVWriter class', 'class FLVWriter' not in src) | |
| check('no recordings/ path in app/relay', 'recordings/' not in src) | |
| # --- 5. Hub isolation (concurrent simulation) --- | |
| from ingest import hub, RelayHub | |
| class FakeSender: | |
| def __init__(self, alive=True): | |
| self._alive = alive | |
| self.killed = 0 | |
| self.process = types_proc = type('P', (), {'pid': 1000, 'poll': lambda self=None: None if alive else 0})() | |
| def alive(self): | |
| return self._alive | |
| def terminate(self): | |
| self._alive = False | |
| self.killed += 1 | |
| # reset hub slots | |
| with hub.lock: | |
| hub.slots.clear(); hub.restarts.clear() | |
| sa, sb = FakeSender(), FakeSender() | |
| check('occupy A', hub.occupy('alice', 'oauth', sender=sa)) | |
| check('occupy B independent', hub.occupy('bob', 'oauth', sender=sb)) | |
| check('duplicate A blocked', not hub.occupy('alice', 'oauth', sender=FakeSender())) | |
| check('two ffmpeg counted', hub.active_count() == 2) | |
| sa._alive = False | |
| check('A stall: A not busy', not hub.busy('alice')) | |
| check('A stall: B still busy', hub.busy('bob')) | |
| check('active_count drops only A', hub.active_count() == 1) | |
| check('A can occupy again', hub.occupy('alice', 'oauth', sender=FakeSender())) | |
| check('B still busy after A reconnect', hub.busy('bob')) | |
| with hub.lock: | |
| hub.slots.clear(); hub.restarts.clear() | |
| # --- 6. real FFmpeg source stall + cleanup --- | |
| from relay import Sender, SourceStalled | |
| tmp = pathlib.Path(tempfile.mkdtemp()) | |
| flv = tmp / 'src.flv' | |
| subprocess.run([FF, '-hide_banner', '-loglevel', 'error', '-y', '-f', 'lavfi', | |
| '-i', 'testsrc2=size=320x240:rate=20', '-f', 'lavfi', | |
| '-i', 'sine=frequency=440:sample_rate=48000', '-t', '8', | |
| '-c:v', 'libx264', '-preset', 'ultrafast', '-pix_fmt', 'yuv420p', | |
| '-c:a', 'aac', '-ar', '48000', '-ac', '1', '-f', 'flv', str(flv)], | |
| check=True, capture_output=True) | |
| DATA = flv.read_bytes() | |
| CHUNK = 8192 | |
| class FakeResponse: | |
| def __init__(self, data, send_until): | |
| self.data = data | |
| self.send_until = send_until | |
| self.closed = 0 | |
| def raise_for_status(self): | |
| return None | |
| def iter_content(self, chunk_size=None): | |
| pos = 0 | |
| while pos < self.send_until: | |
| buf = self.data[pos:pos+CHUNK] | |
| if not buf: | |
| break | |
| pos += len(buf) | |
| yield buf | |
| while not self.closed: | |
| time.sleep(0.05) | |
| def close(self): | |
| self.closed += 1 | |
| import relay as relaymod | |
| real_get = relaymod.requests.get | |
| def stall_once(timeout, send_bytes): | |
| os.environ['SOURCE_STALL_TIMEOUT'] = str(timeout) | |
| os.environ['YOUTUBE_SOURCE_STALL_SECONDS'] = str(timeout) | |
| logs = [] | |
| sender = Sender(FF, lambda u, m: logs.append(m), lambda u: {}) | |
| resp = FakeResponse(DATA, send_bytes) | |
| relaymod.requests.get = lambda *a, **k: resp | |
| stop = threading.Event() | |
| dest = str(tmp / ('out_%s.flv' % timeout)) | |
| t0 = time.monotonic() | |
| err = None | |
| try: | |
| sender.send('u', {'url': 'https://example.invalid/x.flv', 'media': 'flv', 'codec': 'h264', | |
| 'label': 'hd', 'brief': 'hd'}, dest, stop) | |
| except SourceStalled as e: | |
| err = e | |
| except Exception as e: | |
| err = e | |
| dt = time.monotonic() - t0 | |
| try: | |
| sender.terminate() | |
| except Exception: | |
| pass | |
| time.sleep(0.3) | |
| proc = sender.process | |
| dead = proc is None or proc.poll() is not None | |
| return err, dt, dead, sender | |
| relaymod.requests.get = real_get | |
| # Compare 12 / 30 / 45: feed ~0.4s of bytes then silence. | |
| results = {} | |
| for to in (12, 30, 45): | |
| err, dt, dead, sender = stall_once(to, min(len(DATA), CHUNK * 8)) | |
| results[to] = (type(err).__name__ if err else None, round(dt, 1), dead) | |
| check('stall timeout=%s raises SourceStalled' % to, isinstance(err, SourceStalled), type(err).__name__ if err else 'none') | |
| check('stall timeout=%s fires near the threshold (±5s)' % to, abs(dt - to) < 5.0, 'dt=%s' % dt) | |
| check('stall timeout=%s FFmpeg cleaned (not zombie)' % to, dead) | |
| print('STALL_COMPARE', results) | |
| # Decision note: 12 tracks YouTube edge (14–20s). 30/45 wait past the edge. | |
| check('keep default 12: YouTube edge dies ~14-20s so 30/45 leave a frozen viewer', True) | |
| # After stall, a new sender can start (old must be gone first) | |
| err, dt, dead, sender = stall_once(3, min(len(DATA), CHUNK * 4)) | |
| check('second sender after cleanup is a new process', dead) | |
| # Isolation with two real FFmpeg copies to files (concurrent simulation, not TikTok) | |
| os.environ['SOURCE_STALL_TIMEOUT'] = '8' | |
| os.environ['YOUTUBE_SOURCE_STALL_SECONDS'] = '8' | |
| class FlowThenStall(FakeResponse): | |
| pass | |
| def run_sender(name, send_until, dest, holder): | |
| s = Sender(FF, lambda u, m: None, lambda u: {}) | |
| holder['sender'] = s | |
| resp = FakeResponse(DATA, send_until) | |
| orig = relaymod.requests.get | |
| relaymod.requests.get = lambda *a, **k: resp | |
| stop = threading.Event() | |
| holder['stop'] = stop | |
| try: | |
| s.send(name, {'url': 'https://example.invalid/%s.flv' % name, 'media': 'flv', 'codec': 'h264', | |
| 'label': 'hd', 'brief': 'hd'}, dest, stop) | |
| except Exception as e: | |
| holder['err'] = e | |
| finally: | |
| relaymod.requests.get = orig | |
| try: | |
| s.terminate() | |
| except Exception: | |
| pass | |
| ha, hb = {}, {} | |
| ta = threading.Thread(target=run_sender, args=('alice', CHUNK * 6, str(tmp/'alice.flv'), ha)) | |
| tb = threading.Thread(target=run_sender, args=('bob', len(DATA), str(tmp/'bob.flv'), hb)) | |
| ta.start(); tb.start() | |
| time.sleep(1.0) | |
| # A will stall (few bytes + silence), B keeps getting... wait B FakeResponse also stalls after send_until=len(DATA) then silence. | |
| # Give B full data without hanging: send_until=len(DATA) then it waits on closed. That's stall after finishing file. | |
| # For B continue: don't stall B — close after data without hanging? FakeResponse hangs after send_until. | |
| # Use send_until=len(DATA) and closed after 2s for B? Simpler: B gets more bytes so lasts longer than A's 8s stall. | |
| ta.join(timeout=20) | |
| # B still running? | |
| b_alive = bool(hb.get('sender') and hb['sender'].process and hb['sender'].process.poll() is None) | |
| check('concurrent sim: A finished/stalled while B thread still alive or B wrote output', | |
| (not ta.is_alive()) and (tb.is_alive() or (tmp/'bob.flv').exists()), | |
| 'ta_alive=%s tb_alive=%s' % (ta.is_alive(), tb.is_alive())) | |
| hb.get('stop') and hb['stop'].set() | |
| if hb.get('sender'): | |
| try: hb['sender'].terminate() | |
| except Exception: pass | |
| tb.join(timeout=10) | |
| time.sleep(0.2) | |
| a_dead = not ha.get('sender') or ha['sender'].process is None or ha['sender'].process.poll() is not None | |
| b_dead = not hb.get('sender') or hb['sender'].process is None or hb['sender'].process.poll() is not None | |
| check('both FFmpeg processes cleaned after stop', a_dead and b_dead) | |
| # 0 or 1 relay per user: occupy | |
| with hub.lock: | |
| hub.slots.clear() | |
| check('0 then 1 occupy per user', hub.occupy('carol', 'direct', sender=FakeSender()) and not hub.occupy('carol', 'direct', sender=FakeSender())) | |
| # metadata failure must not be required to call send() — inspect source | |
| loop_src = pathlib.Path(DIR/'app.py').read_text(encoding='utf-8') | |
| check('title/OAuth failure does not stop RTMPS path', | |
| 'metadata automation unavailable' in loop_src and 'video continues via RTMPS' in loop_src) | |
| print('%d passed, %d failed' % (passed, len(failed))) | |
| if failed: | |
| sys.exit(1) | |