42 lines
2.8 KiB
Python
42 lines
2.8 KiB
Python
import io
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from unittest.mock import Mock
|
|
from deploy.docker_helper import video_http
|
|
|
|
class Upstream(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def do_GET(self):self.answer()
|
|
def do_POST(self):self.answer()
|
|
def do_DELETE(self):self.answer()
|
|
def answer(self):
|
|
self.server.received=(self.command,self.path,dict(self.headers),self.rfile.read(int(self.headers.get('Content-Length','0'))))
|
|
body=b'\x00original-LTX-response\xff'
|
|
self.send_response(422 if self.path=='/api/invalid' else 206)
|
|
self.send_header('Content-Type','application/octet-stream');self.send_header('Content-Length',str(len(body)));self.send_header('Content-Range','bytes 0-22/23');self.end_headers();self.wfile.write(body)
|
|
class VideoProxyTests(unittest.TestCase):
|
|
def test_preserves_routes_payload_status_range_and_backend_auth(self):
|
|
server=ThreadingHTTPServer(('127.0.0.1',0),Upstream);threading.Thread(target=server.serve_forever,daemon=True).start()
|
|
try:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
token=Path(d)/'token';token.write_text('synthetic-upstream-secret')
|
|
manager=Mock();manager.video_configs.return_value=[dict(id='ltx',api_url=f'http://127.0.0.1:{server.server_port}',token_file=str(token))];manager.video_status.return_value={'running':True}
|
|
for method,path in [('POST','/api/generate?native=1'),('POST','/api/generate/cancel'),('GET','/api/invalid'),('DELETE','/api/models/delete')]:
|
|
payload=b'{"native_field":"unchanged"}';out=io.BytesIO()
|
|
video_http(manager,dict(action='video-http',service='ltx',method=method,path=path,length=len(payload),headers={'Content-Type':'application/json','Range':'bytes=0-22'}),io.BytesIO(payload),out)
|
|
raw=out.getvalue();head,body=raw.split(b'\r\n\r\n',1)
|
|
self.assertIn(b'422' if path.endswith('invalid') else b'206',head)
|
|
self.assertEqual(body,b'\x00original-LTX-response\xff')
|
|
received=server.received;self.assertEqual(received[:2],(method,path));self.assertEqual(received[3],payload)
|
|
self.assertEqual(received[2]['Authorization'],'Bearer synthetic-upstream-secret')
|
|
self.assertNotIn(b'synthetic-upstream-secret',raw)
|
|
finally:server.shutdown();server.server_close()
|
|
def test_rejects_unregistered_targets_header_injection_and_absolute_urls(self):
|
|
manager=Mock();manager.video_configs.return_value=[]
|
|
for changes in ({},{'path':'http://example.com/'},{'headers':{'Authorization':'bad'}},{'headers':{'Accept':'ok\r\nbad'}},{'length':-1}):
|
|
request=dict(action='video-http',service='unknown',method='GET',path='/health',length=0,headers={});request.update(changes);out=io.BytesIO()
|
|
video_http(manager,request,io.BytesIO(),out);self.assertIn(b'502',out.getvalue())
|