Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Regression for live source handover; no real TikTok/YouTube requests. | |
| Starting a replacement pump while swap_source=True made the new thread exit before | |
| reading a single byte. FFmpeg was kept alive but YouTube received no replacement media. | |
| """ | |
| import io | |
| import http.server | |
| import json | |
| import shutil | |
| import subprocess | |
| import tempfile | |
| import os | |
| import pathlib | |
| import sys | |
| import threading | |
| import time | |
| import unittest | |
| from unittest.mock import Mock, patch | |
| ROOT = pathlib.Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else pathlib.Path(__file__).parent | |
| sys.path.insert(0, str(ROOT)) | |
| os.environ.setdefault('DATA_DIR', tempfile.mkdtemp(prefix='handover-test-data-')) | |
| from relay import Sender | |
| from youtube_manager import ManagedYouTube | |
| class InlineThread: | |
| """Deterministic scheduler: the new pump runs immediately when start() is called.""" | |
| def __init__(self, target, args=(), daemon=False): | |
| self.target = target | |
| self.args = args | |
| self.daemon = daemon | |
| self.started = False | |
| def start(self): | |
| self.started = True | |
| self.target(*self.args) | |
| def join(self, timeout=None): | |
| pass | |
| def is_alive(self): | |
| return False | |
| class Response: | |
| def __init__(self, chunks=(b'replacement media',)): | |
| self.chunks = chunks | |
| self.closed = False | |
| self.reads = 0 | |
| def raise_for_status(self): | |
| pass | |
| def iter_content(self, chunk_size=None): | |
| for chunk in self.chunks: | |
| self.reads += 1 | |
| yield chunk | |
| def close(self): | |
| self.closed = True | |
| class HandoverTests(unittest.TestCase): | |
| def setUp(self): | |
| self.env = patch.dict(os.environ, {'SOURCE_STALL_TIMEOUT': '0'}) | |
| self.env.start() | |
| self.addCleanup(self.env.stop) | |
| self.sender = Sender('/bin/true', lambda *a: None) | |
| self.p = Mock() | |
| self.p.poll.return_value = None | |
| self.p.stdin = io.BytesIO() | |
| self.old_pump = Mock() | |
| self.old_pump.is_alive.return_value = False | |
| self.old_response = Response() | |
| self.new_response = Response() | |
| self.stop = threading.Event() | |
| self.sess = {'last': time.monotonic() - 4, 'max_seconds': 180, | |
| 'force_stop': False, 'timeouts': 0, 'error': '', 'ended': True} | |
| self.bridge = Mock() | |
| self.bridge.feed.side_effect = lambda chunk: chunk | |
| self.refresh = Mock(return_value={'url': 'https://example.invalid/source.flv', 'codec': 'h264', 'label': 'hd'}) | |
| def attach(self, refresh=None): | |
| with patch('relay.requests.get', return_value=self.new_response), patch('relay.threading.Thread', InlineThread): | |
| return self.sender._try_reattach('alice', self.p, self.bridge, self.sess, | |
| self.stop, refresh or self.refresh, 10, 3, self.old_response, self.old_pump) | |
| def test_new_pump_reads_media_even_when_scheduled_immediately(self): | |
| attached = self.attach() | |
| self.assertIsNotNone(attached) | |
| self.assertEqual(self.new_response.reads, 1) | |
| self.assertEqual(self.p.stdin.getvalue(), b'replacement media') | |
| self.assertEqual(self.sender.snapshot()['source_bytes'], len(b'replacement media')) | |
| self.assertFalse(self.sess['swap_source']) | |
| def test_old_pump_must_stop_before_bridge_is_reused(self): | |
| self.old_pump.is_alive.return_value = True | |
| self.assertIsNone(self.attach()) | |
| self.bridge.reconnect.assert_not_called() | |
| self.refresh.assert_not_called() | |
| self.assertEqual(self.new_response.reads, 0) | |
| def test_stop_during_resolution_does_not_resurrect_the_reader(self): | |
| def stop_then_refresh(): | |
| self.stop.set() | |
| return self.refresh() | |
| self.assertIsNone(self.attach(stop_then_refresh)) | |
| self.bridge.reconnect.assert_not_called() | |
| self.assertEqual(self.new_response.reads, 0) | |
| def test_refresh_does_not_fabricate_source_heartbeat(self): | |
| self.new_response = Response(chunks=()) | |
| previous = self.sess['last'] | |
| self.assertIsNotNone(self.attach()) | |
| self.assertEqual(self.sess['last'], previous) | |
| self.assertEqual(self.sender.snapshot()['source_bytes'], 0) | |
| def test_opened_source_has_bounded_first_byte_grace(self): | |
| previous = self.sess['last'] | |
| now = time.monotonic() | |
| self.new_response = Response(chunks=()) | |
| self.attach() | |
| self.assertGreater(self.sess['reattach_not_before'], now) | |
| self.assertLessEqual(self.sess['reattach_not_before'], time.monotonic() + 5) | |
| self.assertEqual(self.sess['last'], previous) | |
| def test_cancelled_reader_cannot_deliver_a_late_chunk(self): | |
| cancel = threading.Event() | |
| response = Response() | |
| def late_chunk(_size=None): | |
| cancel.set() | |
| yield b'old reader late packet' | |
| response.iter_content = late_chunk | |
| with patch('relay.threading.Thread', InlineThread): | |
| self.sender._pump_source(response, self.p, self.bridge, self.stop, | |
| self.sess, 10, 3, cancel) | |
| self.assertEqual(self.p.stdin.getvalue(), b'') | |
| self.assertEqual(self.sender.snapshot()['source_bytes'], 0) | |
| def test_no_source_is_a_failed_handover_not_a_fake_success(self): | |
| self.refresh.return_value = None | |
| self.assertIsNone(self.attach()) | |
| self.bridge.reconnect.assert_not_called() | |
| class CanvasTests(unittest.TestCase): | |
| def test_placeholder_16x16_does_not_pin_the_live_canvas(self): | |
| sender = Sender('/bin/true', lambda *a: None) | |
| self.assertEqual(sender.lock_gold({'width': 16, 'height': 16}), (0, 0)) | |
| self.assertEqual(sender.lock_gold({'width': 720, 'height': 1440}), (720, 1280)) | |
| def test_unknown_or_placeholder_matchgop_uses_safe_constant_canvas(self): | |
| for candidate in ({}, {'width': 16, 'height': 16}): | |
| sender = Sender('/bin/true', lambda *a: None) | |
| sender.mode_override = 'matchgop' | |
| args = sender.command('test.flv', candidate) | |
| vf = args[args.index('-vf') + 1] | |
| self.assertIn('scale=w=720:h=1280:', vf) | |
| # A later real configuration does not mutate the encoder canvas. | |
| sender.lock_gold({'width': 960, 'height': 1920}) | |
| self.assertEqual((sender._gop_w, sender._gop_h), (720, 1280)) | |
| def test_normal_and_real_1080_source_dimensions_remain_supported(self): | |
| for size in ((160, 120), (320, 240), (1080, 1920)): | |
| sender = Sender('/bin/true', lambda *a: None) | |
| self.assertEqual(sender.lock_gold({'width': size[0], 'height': size[1]}), size) | |
| class CachedSourceTests(unittest.TestCase): | |
| def setUp(self): | |
| self.tmp = tempfile.TemporaryDirectory() | |
| self.addCleanup(self.tmp.cleanup) | |
| self.manager = ManagedYouTube('/bin/true', lambda u, force=False: None, | |
| lambda *a: None, lambda: ['alice'], pathlib.Path(self.tmp.name) / 'journal.json') | |
| self.manager.video_sources = lambda sources: list(sources) | |
| self.hd = {'url': 'https://example.invalid/hd.flv', 'media': 'flv', 'codec': 'h264', | |
| 'quality_name': 'hd', 'width': 720, 'height': 1280, 'label': 'hd'} | |
| self.sd = dict(self.hd, url='https://example.invalid/sd.flv', quality_name='sd', | |
| width=480, height=854, label='sd') | |
| self.obs = {'status': 'live', 'room_id': 'same-room', 'sources': [self.hd, self.sd]} | |
| self.manager._observe = Mock(return_value=self.obs) | |
| def test_recent_compatible_source_is_reused_without_blocking_metadata_lookup(self): | |
| self.manager._held['alice'] = (time.monotonic(), self.obs) | |
| chosen = self.manager._refresh_flv('alice', self.hd) | |
| self.assertEqual(chosen['url'], self.sd['url']) | |
| self.manager._observe.assert_not_called() | |
| def test_stale_source_snapshot_requires_fresh_lookup(self): | |
| self.manager._held['alice'] = (time.monotonic() - 60, self.obs) | |
| self.assertEqual(self.manager._refresh_flv('alice', self.hd)['url'], self.sd['url']) | |
| self.manager._observe.assert_called_once_with('alice', force=True, bypass_throttle=True, allow_hold=False) | |
| def test_cache_with_wrong_codec_cannot_be_spliced_into_open_publisher(self): | |
| self.manager._held['alice'] = (time.monotonic(), dict(self.obs, sources=[dict(self.sd, codec='h265')])) | |
| chosen = self.manager._refresh_flv('alice', self.hd) | |
| self.assertEqual(chosen['codec'], 'h264') | |
| self.manager._observe.assert_called_once() | |
| class RealHTTPContinuityTests(unittest.TestCase): | |
| def run_continuity(self, mode): | |
| with tempfile.TemporaryDirectory(prefix='source-handover-media-') as tmp: | |
| root = pathlib.Path(tmp) | |
| source = root / 'generated.flv' | |
| output = root / 'output.flv' | |
| subprocess.run([shutil.which('ffmpeg'), '-v', 'error', '-y', '-f', 'lavfi', | |
| '-i', 'testsrc2=size=160x120:rate=20', '-f', 'lavfi', '-i', 'sine=sample_rate=48000', | |
| '-t', '2', '-c:v', 'libx264', '-preset', 'ultrafast', '-pix_fmt', 'yuv420p', | |
| '-g', '20', '-c:a', 'aac', '-f', 'flv', str(source)], check=True, capture_output=True) | |
| media = source.read_bytes() | |
| calls = [] | |
| class Handler(http.server.BaseHTTPRequestHandler): | |
| def log_message(self, *a): | |
| pass | |
| def do_GET(self): | |
| calls.append(time.monotonic()) | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'video/x-flv') | |
| self.send_header('Content-Length', str(len(media))) | |
| self.end_headers() | |
| # Header arrives immediately; replacements wait before their first byte. | |
| # Without startup grace the supervisor closes them before this media arrives. | |
| if len(calls) > 1: | |
| time.sleep(1) | |
| try: | |
| self.wfile.write(media) | |
| self.wfile.flush() | |
| except (OSError, BrokenPipeError): | |
| pass | |
| server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), Handler) | |
| server.daemon_threads = True | |
| server_thread = threading.Thread(target=server.serve_forever, daemon=True) | |
| server_thread.start() | |
| stop = threading.Event() | |
| errors = [] | |
| candidate = {'url': 'http://127.0.0.1:%d/source.flv' % server.server_port, | |
| 'media': 'flv', 'codec': 'h264', 'label': 'fixture', 'width': 160, 'height': 120} | |
| sender = Sender(shutil.which('ffmpeg'), lambda *a: None) | |
| sender.mode_override = mode | |
| spawned = [] | |
| from relay import subprocess as relay_subprocess | |
| real_popen = relay_subprocess.Popen | |
| def track_process(*a, **kw): | |
| child = real_popen(*a, **kw) | |
| spawned.append(child) | |
| return child | |
| def send(): | |
| try: | |
| sender.send('fixture', candidate, str(output), stop, refresh=lambda: dict(candidate)) | |
| except Exception as exc: | |
| errors.append(exc) | |
| env = {'SOURCE_STALL_TIMEOUT': '3', 'YOUTUBE_REATTACH_GRACE_SECONDS': '2', | |
| 'YOUTUBE_SOURCE_READ_TIMEOUT': '5', 'YOUTUBE_PUMP_JOIN_SECONDS': '2', | |
| 'YOUTUBE_SEND_MODE': mode, 'YOUTUBE_VIDEO_MODE': 'original', | |
| 'YOUTUBE_AUTO_TRANSCODE': '0', 'YOUTUBE_X264_PRESET': 'ultrafast'} | |
| thread = None | |
| try: | |
| with patch.dict(os.environ, env), patch('relay.subprocess.Popen', side_effect=track_process): | |
| thread = threading.Thread(target=send, daemon=True) | |
| thread.start() | |
| time.sleep(7.5) | |
| alive = thread.is_alive() and sender.alive() | |
| snapshot = sender.snapshot() | |
| stop.set() | |
| thread.join(6) | |
| self.assertTrue(alive, 'Publisher ended instead of reusing the delayed replacement source: %r' % errors) | |
| self.assertFalse(errors, repr(errors)) | |
| self.assertEqual(len(spawned), 1, 'Source reconnect must keep the same FFmpeg process') | |
| self.assertGreaterEqual(len(calls), 3) | |
| self.assertGreater(snapshot.get('output_seconds', 0), 4) | |
| self.assertFalse(thread.is_alive(), 'Sender did not stop cleanly') | |
| probe = subprocess.run([shutil.which('ffprobe'), '-v', 'error', '-show_entries', | |
| 'format=duration', '-of', 'json', str(output)], check=True, capture_output=True, text=True) | |
| self.assertGreater(float(json.loads(probe.stdout)['format']['duration']), 4) | |
| finally: | |
| stop.set() | |
| sender.terminate() | |
| if thread: | |
| thread.join(5) | |
| server.shutdown() | |
| server.server_close() | |
| def test_eof_and_delayed_replacements_keep_same_copy_publisher(self): | |
| self.run_continuity('copy') | |
| def test_eof_and_delayed_replacements_keep_same_matchgop_publisher(self): | |
| self.run_continuity('matchgop') | |
| if __name__ == '__main__': | |
| unittest.main(argv=[sys.argv[0]], verbosity=2) | |