import importlib.util import json import os import subprocess import tempfile import time import unittest from pathlib import Path ROOT = Path(__file__).resolve().parents[1] spec = importlib.util.spec_from_file_location('worker', ROOT / 'pt-worker.py') worker = importlib.util.module_from_spec(spec) spec.loader.exec_module(worker) class PolicyTests(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.old_state = worker.STATE worker.STATE = Path(self.tmp.name) self.meter = worker.Meter() self.now = time.time() for direction in ('download', 'upload'): self.meter.consume({'kind': 'transfer', 'time': self.now, 'ok': True, 'direction': direction, 'bytes': 10_000_000, 'seconds': 1}) self.policy = worker.Admission({}) self.sample = {'tailscale_path': 'direct', 'latency_ms': 5, 'network_ok': True, 'registered': True, 'qsv': {'qsv_h264_ok': True, 'tested_at': self.now}, 'gpu_ready': True, 'gpu_tested_at': self.now, 'gpu_accessible': True, 'on_ac': True, 'disk_free_bytes': 100 * 1024**3, 'temperature_c': 50, 'throttle_counts': {}, 'active_jobs': 0} self.control = {'paused': False} def tearDown(self): worker.STATE = self.old_state self.tmp.cleanup() def decide(self): return self.policy.decide(self.sample, self.meter, self.control, self.now) def available(self): self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.assertEqual(self.decide(), 'AVAILABLE') def test_recovery_hysteresis_and_low_bandwidth_drain(self): self.available() self.meter.data['download']['mbps'] = 1 self.assertEqual(self.decide(), 'AVAILABLE') self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.meter.data['download']['mbps'] = 11 # insufficient recovery margin self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.meter.data['download']['mbps'] = 20 self.available() def test_path_transition_invalidates_real_bandwidth(self): self.available() self.sample['tailscale_path'] = 'DERP' self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.assertIsNone(self.sample['recent_download_mbps']) self.sample['tailscale_path'] = 'direct' self.assertEqual(self.decide(), 'NETWORK_DEGRADED') def test_battery_gpu_disk_and_network_refuse(self): self.available() self.sample['on_ac'] = False self.assertEqual(self.decide(), 'DRAINING') self.sample['on_ac'] = True self.sample['gpu_accessible'] = False self.assertEqual(self.decide(), 'GPU_UNAVAILABLE') self.sample['gpu_accessible'] = True self.sample['disk_free_bytes'] = 0 self.assertEqual(self.decide(), 'LOW_DISK') self.sample['disk_free_bytes'] = 100 * 1024**3 self.sample['network_ok'] = False self.assertEqual(self.decide(), 'OFFLINE') def test_gpu_failure_latched_until_later_selftest(self): self.meter.consume({'kind': 'gpu_failure', 'time': self.now + 1}) self.assertEqual(self.decide(), 'GPU_UNAVAILABLE') self.sample['qsv']['tested_at'] = self.now + 2 self.sample['gpu_tested_at'] = self.now + 2 self.available() def test_transfer_failure_cooldown_and_stale_metrics(self): self.meter.consume({'kind': 'transfer', 'time': self.now, 'ok': False}) self.assertEqual(self.decide(), 'NETWORK_DEGRADED') self.now += 61 self.available() self.now += 900 self.assertEqual(self.decide(), 'AVAILABLE') self.assertEqual(self.decide(), 'NETWORK_DEGRADED') def test_thermal_throttle_and_pause(self): self.available() self.sample['throttle_counts'] = {'cpu0': 1} self.decide() self.sample['throttle_counts'] = {'cpu0': 2} self.assertEqual(self.decide(), 'THERMAL_THROTTLED') self.sample['temperature_c'] = 50 self.now += 121 self.available() self.control['paused'] = True self.assertEqual(self.decide(), 'PAUSED') def test_job_correlation_uses_uuid_under_concurrency(self): for uuid in ('a', 'b'): self.meter.consume({'kind': 'job_accept', 'time': self.now, 'ok': True, 'uuid': uuid}) self.meter.consume({'kind': 'encode', 'time': self.now + 2, 'job_uuid': 'b', 'duration': 200}) self.meter.consume({'kind': 'encode', 'time': self.now + 3, 'job_uuid': 'a', 'duration': 100}) self.meter.consume({'kind': 'job_success', 'time': self.now + 10, 'uuid': 'b', 'ok': True}) self.assertEqual(self.meter.data['last_job_realtime_factor'], 20) class AdapterTests(unittest.TestCase): def probe(self, codec='h264', rotation=False): return {'streams': [{'codec_type': 'video', 'codec_name': codec, 'pix_fmt': 'yuv420p', 'width': 1920, 'height': 1080, 'side_data_list': [{'rotation': 90}] if rotation else []}], 'format': {'duration': '10'}} def test_qsv_decode_scale_encode_and_no_x264_flags(self): args, info = worker.qsv.rewrite(['-i', '/tmp/in', '-vcodec', 'libx264', '-vf', 'scale=w=-2:h=720', '-pix_fmt', 'yuv420p', '-preset', 'veryfast', '-b_strategy', '1', '-bf', '16', '/tmp/out'], self.probe()) self.assertIn('h264_qsv', args) self.assertIn('scale_qsv=w=1280:h=720:format=nv12', args) self.assertNotIn('libx264', args) self.assertNotIn('veryfast', args) self.assertNotIn('-b_strategy', args) self.assertEqual(info['target_height'], 720) def test_unknown_decode_rotation_and_cpu_encoder_fail(self): for probe in (self.probe('mpeg4'), self.probe(rotation=True)): with self.assertRaises(ValueError): worker.qsv.rewrite(['-i', '/tmp/in', '-c:v', 'libx264', '/tmp/out'], probe) with self.assertRaises(ValueError): worker.qsv.rewrite(['-i', '/tmp/in', '-c:v', 'libx265', '/tmp/out'], self.probe()) def test_unreviewed_software_filter_refused(self): with self.assertRaises(ValueError): worker.qsv.rewrite(['-i', '/tmp/in', '-c:v', 'libx264', '-vf', 'scale=640:360,eq=contrast=2', '/tmp/out'], self.probe()) def test_real_missing_qsv_refuses_without_cpu_output(self): with tempfile.TemporaryDirectory() as tmp: src, out = Path(tmp) / 'src.mp4', Path(tmp) / 'out.mp4' subprocess.run(['/usr/bin/ffmpeg', '-v', 'error', '-f', 'lavfi', '-i', 'testsrc2=size=64x64:rate=2', '-t', '1', '-c:v', 'libx264', '-threads', '1', str(src)], check=True) env = {**os.environ, 'PT_STATE_DIR': tmp, 'PT_RENDER_DEVICE': tmp + '/missing-device'} p = subprocess.run(['/usr/bin/python3', str(ROOT / 'ffmpeg-qsv.py'), '-i', str(src), '-c:v', 'libx264', str(out)], env=env, capture_output=True) self.assertEqual(p.returncode, 78) self.assertFalse(out.exists()) self.assertIn('gpu_failure', (Path(tmp) / 'events.jsonl').read_text()) if __name__ == '__main__': unittest.main()