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()