24 lines
1.4 KiB
Python
24 lines
1.4 KiB
Python
import http.client
|
|
import threading
|
|
import unittest
|
|
from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer
|
|
from video_comfy_proxy import relay
|
|
from test_video_proxy import Upstream
|
|
|
|
class Proxy(BaseHTTPRequestHandler):
|
|
def log_message(self,*args):pass
|
|
def do_POST(self):relay(self,self.server.target)
|
|
def do_GET(self):relay(self,self.server.target)
|
|
|
|
class ComfyProxyTests(unittest.TestCase):
|
|
def test_binary_ranges_and_query_preserved_credentials_stripped(self):
|
|
upstream=ThreadingHTTPServer(('127.0.0.1',0),Upstream);proxy=ThreadingHTTPServer(('127.0.0.1',0),Proxy);proxy.target=upstream.server_port
|
|
for server in (upstream,proxy):threading.Thread(target=server.serve_forever,daemon=True).start()
|
|
try:
|
|
connection=http.client.HTTPConnection('127.0.0.1',proxy.server_port)
|
|
connection.request('POST','/prompt?native=1',body=b'synthetic',headers={'Authorization':'Bearer synthetic-secret','Cookie':'deck_session=synthetic','Content-Type':'application/json','Range':'bytes=0-22'})
|
|
response=connection.getresponse();self.assertEqual(response.status,206);self.assertEqual(response.read(),b'\x00original-LTX-response\xff');connection.close()
|
|
received=upstream.received;self.assertEqual(received[:2],('POST','/prompt?native=1'));self.assertEqual(received[3],b'synthetic');self.assertNotIn('Cookie',received[2]);self.assertNotIn('Authorization',received[2])
|
|
finally:
|
|
for server in (proxy,upstream):server.shutdown();server.server_close()
|